Filter: 100%
 50000/50000 [00:15<00:00, 6263.21 examples/s]
Initialized a new FAISS index.
Index: -1, Similarity Score: 0.00%
Index: -1, Similarity Score: 0.00%
Index: -1, Similarity Score: 0.00%
Index: -1, Similarity Score: 0.00%
Index: -1, Similarity Score: 0.00%

above results are wrong , fix the code make more robustness , generalised
import os
from typing import List, Tuple, Any

import torch
from PIL import Image
from transformers import AutoProcessor, CLIPModel, AutoTokenizer
from datasets import load_dataset, Dataset
import faiss
import numpy as np
from tqdm import tqdm
import plotly.graph_objects as go
import plotly.io as pio


class CLIPSimilaritySearch:
    """
    A class to perform prompt-based image similarity search using CLIP and FAISS,
    with results visualized using Plotly.
    """

    def __init__(
        self,
        model_name: str = "openai/clip-vit-base-patch32",
        dataset_name: str = "cifar10",
        dataset_split: str = "train",
        max_images: int = 10000,
        index_path: str = "clip.index",
        device: str = None,
    ) -> None:
        """
        Initialize the CLIPSimilaritySearch pipeline.

        Args:
            model_name (str): Name of the pre-trained CLIP model.
            dataset_name (str): Name of the dataset to load.
            dataset_split (str): Which split of the dataset to use.
            max_images (int): Maximum number of images to process.
            index_path (str): Path to save/load the FAISS index.
            device (str, optional): Device to run the model on. Defaults to CUDA if available.
        """
        self.model_name = model_name
        self.dataset_name = dataset_name
        self.dataset_split = dataset_split
        self.max_images = max_images
        self.index_path = index_path
        self.device = torch.device(device) if device else torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )

        self._load_model()
        self._load_tokenizer()
        self._load_dataset()
        self._initialize_faiss_index()

    def _load_model(self) -> None:
        """Load the CLIP model and processor."""
        try:
            self.processor = AutoProcessor.from_pretrained(self.model_name)
            self.model = CLIPModel.from_pretrained(self.model_name).to(self.device)
        except Exception as e:
            raise RuntimeError(f"Error loading model '{self.model_name}': {e}")

    def _load_tokenizer(self) -> None:
        """Load the tokenizer associated with the CLIP model."""
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        except Exception as e:
            raise RuntimeError(f"Error loading tokenizer for '{self.model_name}': {e}")

    def _load_dataset(self) -> None:
        """Load and preprocess the dataset."""
        try:
            dataset = load_dataset(self.dataset_name)
            self.train_dataset: Dataset = dataset[self.dataset_split].filter(
                lambda example, idx: idx < self.max_images, with_indices=True
            )
            if len(self.train_dataset) == 0:
                raise ValueError("The filtered dataset is empty.")
        except Exception as e:
            raise RuntimeError(f"Error loading dataset '{self.dataset_name}': {e}")

    def _initialize_faiss_index(self) -> None:
        """Initialize or load the FAISS index."""
        embedding_dim = self.model.visual_projection.out_features  # Typically 512 for CLIP
        if os.path.exists(self.index_path):
            try:
                self.index = faiss.read_index(self.index_path)
                if self.index.d == embedding_dim:
                    print("FAISS index loaded successfully.")
                else:
                    raise ValueError("Dimension mismatch in FAISS index.")
            except Exception as e:
                raise RuntimeError(f"Error loading FAISS index from '{self.index_path}': {e}")
        else:
            self.index = faiss.IndexFlatL2(embedding_dim)
            print("Initialized a new FAISS index.")

    def _add_vectors_to_index(self, embeddings: np.ndarray) -> None:
        """Add normalized embeddings to the FAISS index."""
        try:
            faiss.normalize_L2(embeddings)
            self.index.add(embeddings)
        except Exception as e:
            raise RuntimeError(f"Error adding vectors to FAISS index: {e}")

    def _extract_image_features(self, images: List[Image.Image]) -> np.ndarray:
        """Extract image features using CLIP."""
        try:
            inputs = self.processor(images=images, return_tensors="pt", padding=True).to(self.device)
            with torch.no_grad():
                image_features = self.model.get_image_features(**inputs)
            image_features = image_features.cpu().numpy().astype(np.float32)
            return image_features
        except Exception as e:
            raise RuntimeError(f"Error extracting image features: {e}")

    def build_index(self, batch_size: int = 32) -> None:
        """
        Build the FAISS index by processing the dataset in batches.

        Args:
            batch_size (int): Number of images to process in each batch.
        """
        try:
            total = len(self.train_dataset)
            for start_idx in tqdm(range(0, total, batch_size), desc="Building FAISS Index"):
                end_idx = min(start_idx + batch_size, total)
                batch_images = [
                    self.train_dataset[i]["img"] for i in range(start_idx, end_idx)
                ]
                embeddings = self._extract_image_features(batch_images)
                self._add_vectors_to_index(embeddings)
            faiss.write_index(self.index, self.index_path)
            print(f"FAISS index built and saved to '{self.index_path}'.")
        except Exception as e:
            raise RuntimeError(f"Error building FAISS index: {e}")

    def search(
        self, prompt: str, top_k: int = 5
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Search for images similar to the given prompt.

        Args:
            prompt (str): The text prompt to search for.
            top_k (int): Number of top similar images to retrieve.

        Returns:
            Tuple[np.ndarray, np.ndarray]: Distances and indices of the top_k results.
        """
        try:
            text_inputs = self.tokenizer(
                [prompt], return_tensors="pt", padding=True
            ).to(self.device)
            with torch.no_grad():
                text_features = self.model.get_text_features(**text_inputs)
            text_features = text_features.cpu().numpy().astype(np.float32)
            faiss.normalize_L2(text_features)
            distances, indices = self.index.search(text_features, top_k)
            return distances, indices
        except Exception as e:
            raise RuntimeError(f"Error during search: {e}")

    def visualize_results(
        self, distances: np.ndarray, indices: np.ndarray, prompt: str, output_html: str
    ) -> None:
        """
        Visualize the search results using Plotly and save as an HTML file.

        Args:
            distances (np.ndarray): Distances returned by FAISS.
            indices (np.ndarray): Indices of the top_k results.
            prompt (str): The original search prompt.
            output_html (str): Path to save the HTML visualization.
        """
        try:
            fig = go.Figure()
            for rank, (dist, idx) in enumerate(zip(distances[0], indices[0]), start=1):
                similarity = 1 / (1 + dist) * 100  # Similarity score between 0 and 100
                image = self.train_dataset[int(idx)]["img"]
                img_path = f"./temp_{idx}.png"
                image.save(img_path)
                fig.add_trace(
                    go.Image(
                        source=img_path,
                        name=f"Rank {rank}: Similarity {similarity:.2f}%",
                    )
                )
                os.remove(img_path)  # Clean up temporary image file

            fig.update_layout(
                title=f"Top {len(indices[0])} Results for '{prompt}'",
                template="plotly_white",
            )
            pio.write_html(fig, file=output_html, auto_open=False)
            print(f"Visualization saved to '{output_html}'.")
        except Exception as e:
            raise RuntimeError(f"Error during visualization: {e}")

    def run_search_pipeline(self, prompt: str, top_k: int = 5, output_html: str = "results.html") -> None:
        """
        Execute the complete search pipeline: build index, perform search, and visualize.

        Args:
            prompt (str): The text prompt to search for.
            top_k (int, optional): Number of top results to retrieve. Defaults to 5.
            output_html (str, optional): Path to save the HTML visualization. Defaults to 'results.html'.
        """
        if not self.index.is_trained:
            print("Building FAISS index...")
            self.build_index()

        distances, indices = self.search(prompt, top_k)
        for dist, idx in zip(distances[0], indices[0]):
            similarity = (1 / (1 + dist)) * 100
            print(f"Index: {idx}, Similarity Score: {similarity:.2f}%")
        self.visualize_results(distances, indices, prompt, output_html)


def main() -> None:
    """
    Main function to execute the CLIP similarity search pipeline.
    """
    try:
        search_pipeline = CLIPSimilaritySearch(
            model_name="openai/clip-vit-base-patch32",
            dataset_name="cifar10",
            dataset_split="train",
            max_images=10000,
            index_path="clip.index",
        )
        search_prompt = "a photo of a dog"
        search_pipeline.run_search_pipeline(
            prompt=search_prompt, top_k=5, output_html="similarity_results.html"
        )
    except Exception as e:
        print(f"An error occurred: {e}")


if __name__ == "__main__":
    main()