Create app.py

app.py

@@ -0,0 +1,65 @@
+import gradio as gr  
+import pickle
+from datasets import load_dataset
+from torch import nn
+import numpy as np
+
+import torch
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel
+from datasets import load_dataset
+
+def get_clip_embeddings(input_data, input_type='text'):
+    # Load the CLIP model and processor
+    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+    # Prepare the input based on the type
+    if input_type == 'text':
+        inputs = processor(text=input_data, return_tensors="pt", padding=True, truncation=True)
+    elif input_type == 'image':
+        if isinstance(input_data, str):
+            image = Image.open(input_data)
+        elif isinstance(input_data, Image.Image):
+            image = input_data
+        else:
+            raise ValueError("For image input, provide either a file path or a PIL Image object")
+        inputs = processor(images=image, return_tensors="pt")
+    else:
+        raise ValueError("Invalid input_type. Choose 'text' or 'image'")
+
+    # Get the embeddings
+    with torch.no_grad():
+        if input_type == 'text':
+            embeddings = model.get_text_features(**inputs)
+        else:
+            embeddings = model.get_image_features(**inputs)
+
+    return embeddings.numpy()
+
+
+veggies = load_dataset('vojtam/vegetables')
+
+
+text = gr.Textbox(label = "Enter the text")
+image = gr.Gallery()
+
+def get_similar_images(text, n = 4):
+    with open('img_embeddings.pkl', 'rb') as file:
+        img_embeddings = pickle.load(file)
+    text_embedding = get_clip_embeddings(text, input_type='text')
+    cos = nn.CosineSimilarity(dim=1, eps=1e-6)
+    sims = cos(torch.tensor(text_embedding), torch.tensor(img_embeddings))
+    top_n = np.argsort(np.array(sims))[::-1][:4]
+    print(top_n)
+    print(img_embeddings)
+    imgs = []
+    
+    for index in top_n:
+        imgs.append(veggies['train'][index.item()]['image'])
+    return imgs
+    
+
+
+intf = gr.Interface(fn = get_similar_images, inputs = text, outputs = image)
+intf.launch(share=True)