RobotJelly/Text_Or_Image-To-Image_Search

1# Import Libraries2from pathlib import Path3import pandas as pd4import numpy as np5import torch6import pickle7from PIL import Image8from io import BytesIO9import requests10import gradio as gr11import os12import sentence_transformers13from sentence_transformers import SentenceTransformer, util14 15# check if CUDA available 16device = "cuda" if torch.cuda.is_available() else "cpu"17 18IMAGES_DIR = Path("photos/")19   20#Load CLIP model21model = SentenceTransformer('clip-ViT-B-32') 22 23# pre-computed embeddings24emb_filename = 'unsplash-25k-photos-embeddings.pkl'25with open(emb_filename, 'rb') as emb:26        img_names, img_emb = pickle.load(emb)27 28def display_matches(similarity, topk):29    best_matched_images = []30    top_k_indices = torch.topk(similarity, topk, 0).indices31    for matched_image in top_k_indices:32        img = Image.open(IMAGES_DIR / img_names[matched_image])33        best_matched_images.append(img)34    return best_matched_images35  36def image_search(Option, topk, search_text, search_image):37  topk = topk+138  # Input Text Query39  if Option == "Text-To-Image" :40    # Encode the given Input  text for Search & take it in tensor form41     text_emb = model.encode([search_text], convert_to_tensor=True)42     # Compute cosine similarities between encoded input text (in tensor) & encoded images from unsplash dataset43     similarity = util.cos_sim(img_emb, text_emb)44     45     #using the computed similarities, find the topk best matches46     return display_matches(similarity, topk)47  elif Option == "Image-To-Image":48      # Encode the given Input Image for Search & take it in tensor form49      image_emb = model.encode([Image.fromarray(search_image)], convert_to_tensor=True)50      # Compute cosine similarities between encoded input image (in tensor) & encoded images from unsplash dataset51      similarity = util.cos_sim(img_emb, image_emb)52      53      #using the computed similarities, find the topk best matches54      return display_matches(similarity, topk)55  56gr.Interface(fn=image_search, title="Search Image",57             description="Enter the text or image to search for the most relevant images...",58             article=""" 59                     Instructions:- 60                      1. Select the option - `Text to Image`  OR `Image To Image`.61                      2. Select the no. of most relevant images you want to see. 62                      3. Then accordingly enter the text or image.63                      4. Then you will get the images on right. To enter another text/image first clear it then follow steps 1-3.64                     """,65             theme="huggingface",66            inputs=[gr.inputs.Dropdown(["Text-To-Image", "Image-To-Image"]),67                    gr.inputs.Dropdown(["1", "2", "3", "4", "5", "6", "7", "8", "9", "10"], type="index", default="1", label="Select Top K Images"),68                    gr.inputs.Textbox(lines=3, label="Input Text", placeholder="Enter the text..."),69                    gr.inputs.Image(optional=True)70                    ], 71            outputs=gr.outputs.Carousel([gr.outputs.Image(type="pil")]),72            enable_queue=True73             ).launch(debug=True,share=True)