engine/search.py

import os
import numpy as np
from typing import List, Tuple
import torch
from glob import glob
from PIL import Image
from tqdm import tqdm
import matplotlib.pyplot as plt
from transformers import CLIPProcessor, CLIPModel
from sentence_transformers import SentenceTransformer
import sqlite3
from .vector_database import (
    VectorDB,
    ImageEmbeddingCollectionSchema,
    TextEmbeddingCollectionSchema,
)


class ImageSearchModule:
    def __init__(
        self,
        image_embeddings_dir: str,
        original_images_dir: str,
        sqlite_db_path: str = "image_tracker.db",
    ):
        self.image_embeddings_dir = image_embeddings_dir
        self.original_images_dir = original_images_dir
        self.vector_db = VectorDB()
        self.vector_db.create_collection(ImageEmbeddingCollectionSchema)
        self.vector_db.create_collection(TextEmbeddingCollectionSchema)

        self.clip_model = CLIPModel.from_pretrained(
            "wkcn/TinyCLIP-ViT-8M-16-Text-3M-YFCC15M"
        )
        self.clip_preprocess = CLIPProcessor.from_pretrained(
            "wkcn/TinyCLIP-ViT-8M-16-Text-3M-YFCC15M"
        )
        self.text_embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

        self.sqlite_conn = sqlite3.connect(sqlite_db_path)
        self._create_sqlite_table()

    def _create_sqlite_table(self):
        cursor = self.sqlite_conn.cursor()
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS added_images (
                image_name TEXT PRIMARY KEY
            )
        """
        )
        self.sqlite_conn.commit()

    def add_images(self):
        print("Adding images to vector databases")
        cursor = self.sqlite_conn.cursor()

        for filename in tqdm(os.listdir(self.image_embeddings_dir)):
            if filename.startswith("resized_") and filename.endswith("_clip.npy"):
                image_name = filename[
                    8:-9
                ]  # Remove "resized_" prefix and "_clip.npy" suffix

                cursor.execute(
                    "SELECT 1 FROM added_images WHERE image_name = ?", (image_name,)
                )
                if cursor.fetchone() is None:
                    clip_embedding_path = os.path.join(
                        self.image_embeddings_dir, filename
                    )
                    caption_embedding_path = os.path.join(
                        self.image_embeddings_dir, f"resized_{image_name}_caption.npy"
                    )

                    if os.path.exists(clip_embedding_path) and os.path.exists(
                        caption_embedding_path
                    ):
                        with open(clip_embedding_path, "rb") as buffer:
                            image_embedding = np.frombuffer(
                                buffer.read(), dtype=np.float32
                            ).reshape(512)
                        with open(caption_embedding_path, "rb") as buffer:
                            text_embedding = np.frombuffer(
                                buffer.read(), dtype=np.float32
                            ).reshape(384)

                        if self.vector_db.insert_record(
                            ImageEmbeddingCollectionSchema.collection_name,
                            image_embedding,
                            image_name,
                        ):
                            self.vector_db.insert_record(
                                TextEmbeddingCollectionSchema.collection_name,
                                text_embedding,
                                image_name,
                            )
                            cursor.execute(
                                "INSERT INTO added_images (image_name) VALUES (?)",
                                (image_name,),
                            )
                            self.sqlite_conn.commit()

        print("Finished adding images to vector databases")

    def search_by_image(
        self, query_image_path: str, top_k: int = 5, similarity_threshold: float = 0.5
    ) -> List[Tuple[str, float]]:
        if not os.path.exists(query_image_path):
            print(f"Image file not found: {query_image_path}")
            return []
        try:
            query_image = Image.open(query_image_path)
            query_embedding = self._get_image_embedding(query_image)
            results = self.vector_db.client.search(
                collection_name=ImageEmbeddingCollectionSchema.collection_name,
                data=[query_embedding],
                output_fields=["filename"],
                search_params={"metric_type": "COSINE"},
                limit=top_k,
            ).pop()
            return [(item["entity"]["filename"], item["distance"]) for item in results if item["distance"] >= similarity_threshold]
        except Exception as e:
            print(f"Error processing image: {e}")
            return []

    def search_by_text(
        self, query_text: str, top_k: int = 5,similarity_threshold: float = 0.5
    ) -> List[Tuple[str, float]]:
        if not query_text.strip():
            print("Empty text query")
            return []
        try:
            query_embedding = self._get_text_embedding(query_text)
            results = self.vector_db.client.search(
                collection_name=TextEmbeddingCollectionSchema.collection_name,
                data=[query_embedding],
                search_params={"metric_type": "COSINE"},
                output_fields=["filename"],
                limit=top_k,
            ).pop()
            return [(item["entity"]["filename"], item["distance"]) for item in results if item["distance"] >= similarity_threshold]
        except Exception as e:
            print(f"Error processing text: {e}")
            return []

    def _get_image_embedding(self, image: Image.Image) -> np.ndarray:
        with torch.no_grad():
            image_input = self.clip_preprocess(images=image, return_tensors="pt")[
                "pixel_values"
            ].to(self.clip_model.device)
            image_features = self.clip_model.get_image_features(image_input)
        return image_features.cpu().numpy().flatten()

    def _get_text_embedding(self, text: str) -> np.ndarray:
        with torch.no_grad():
            embedding = self.text_embedding_model.encode(text).flatten()
        return embedding

    def display_results(self, results: List[Tuple[str, float]]):
        if not results:
            print("No results to display.")
            return

        num_images = min(5, len(results))
        fig, axes = plt.subplots(1, num_images, figsize=(20, 4))
        axes = [axes] if num_images == 1 else axes

        for i, (image_name, similarity) in enumerate(results[:num_images]):
            pattern = os.path.join(
                self.original_images_dir, f"resized_{image_name}" + "*"
            )
            matching_files = glob(pattern)
            if matching_files:
                image_path = matching_files[0]
                img = Image.open(image_path)
                axes[i].imshow(img)
                axes[i].set_title(f"Similarity: {similarity:.2f}")
                axes[i].axis("off")
            else:
                print(f"No matching image found for {image_name}")
                axes[i].text(0.5, 0.5, "Image not found", ha="center", va="center")
                axes[i].axis("off")

        plt.tight_layout()
        plt.show()

    def __del__(self):
        if hasattr(self, "sqlite_conn"):
            self.sqlite_conn.close()


if __name__ == "__main__":
    from pathlib import Path
    import requests

    PROJECT_ROOT = Path(__file__).resolve().parent.parent
    search = ImageSearchModule(
        image_embeddings_dir=str(PROJECT_ROOT / "data/features"),
        original_images_dir=str(PROJECT_ROOT / "data/images"),
    )
    search.add_images()

    # Search by image
    img_url = (
        "https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg"
    )
    raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
    raw_image.save(PROJECT_ROOT / "test.jpg")
    image_results = search.search_by_image(str(PROJECT_ROOT / "test.jpg"))
    print("Image search results:")
    search.display_results(image_results)

    # Search by text
    text_results = search.search_by_text("Images of Nature")
    print("Text search results:")
    search.display_results(text_results)