Reference for ultralytics/solutions/similarity_search.py

def __call__(self, query):
    """Direct call for search function."""
    return self.search(query)

def extract_image_feature(self, path):
    """Extract CLIP image embedding."""
    image = Image.open(path)
    tensor = self.preprocess(image).unsqueeze(0).to(self.device)
    with torch.no_grad():
        return self.model.encode_image(tensor).cpu().numpy()

def extract_text_feature(self, text):
    """Extract CLIP text embedding."""
    tokens = self.clip.tokenize([text]).to(self.device)
    with torch.no_grad():
        return self.model.encode_text(tokens).cpu().numpy()

def load_or_build_index(self):
    """Loads FAISS index or builds a new one from image features."""
    # Check if the FAISS index and corresponding image paths already exist
    if Path(self.faiss_index).exists() and Path(self.data_path_npy).exists():
        self.LOGGER.info("Loading existing FAISS index...")
        self.index = self.faiss.read_index(self.faiss_index)  # Load the FAISS index from disk
        self.image_paths = np.load(self.data_path_npy)  # Load the saved image path list
        return  # Exit the function as the index is successfully loaded

    # If the index doesn't exist, start building it from scratch
    self.LOGGER.info("Building FAISS index from images...")
    vectors = []  # List to store feature vectors of images

    # Iterate over all image files in the data directory
    for file in self.data_dir.iterdir():
        # Skip files that are not valid image formats
        if file.suffix.lower().lstrip(".") not in IMG_FORMATS:
            continue
        try:
            # Extract feature vector for the image and add to the list
            vectors.append(self.extract_image_feature(file))
            self.image_paths.append(file.name)  # Store the corresponding image name
        except Exception as e:
            self.LOGGER.warning(f"Skipping {file.name}: {e}")

    # If no vectors were successfully created, raise an error
    if not vectors:
        raise RuntimeError("No image embeddings could be generated.")

    vectors = np.vstack(vectors).astype("float32")  # Stack all vectors into a NumPy array and convert to float32
    self.faiss.normalize_L2(vectors)  # Normalize vectors to unit length for cosine similarity

    self.index = self.faiss.IndexFlatIP(vectors.shape[1])  # Create a new FAISS index using inner product
    self.index.add(vectors)  # Add the normalized vectors to the FAISS index
    self.faiss.write_index(self.index, self.faiss_index)  # Save the newly built FAISS index to disk
    np.save(self.data_path_npy, np.array(self.image_paths))  # Save the list of image paths to disk

    self.LOGGER.info(f"Indexed {len(self.image_paths)} images.")

def search(self, query, k=30, similarity_thresh=0.1):
    """Returns top-k semantically similar images to the given query."""
    text_feat = self.extract_text_feature(query).astype("float32")
    self.faiss.normalize_L2(text_feat)

    D, index = self.index.search(text_feat, k)
    results = [
        (self.image_paths[i], float(D[0][idx])) for idx, i in enumerate(index[0]) if D[0][idx] >= similarity_thresh
    ]
    results.sort(key=lambda x: x[1], reverse=True)

    self.LOGGER.info("\nRanked Results:")
    for name, score in results:
        self.LOGGER.info(f"  - {name} | Similarity: {score:.4f}")

    return [r[0] for r in results]

def index(self):
    """Function to process the user query and display output."""
    results = []
    if self.request.method == "POST":
        query = self.request.form.get("query", "").strip()
        results = self.searcher(query)
    return self.render_template("similarity-search.html", results=results)

def run(self, debug=False):
    """Runs the Flask web app."""
    self.app.run(debug=debug)