Upload 2 files

app.py

@@ -0,0 +1,78 @@
+import gradio as gr
+import os
+from datasets import load_dataset
+import torch
+import clip
+from PIL import Image
+import pyarrow as pa
+import lancedb
+
+ds = load_dataset("vipulmaheshwari/GTA-Image-Captioning-Dataset")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, preprocess = clip.load("ViT-L/14", device=device)
+
+
+#Embedding Image Function
+def embedding_image(image):
+  processed_image = preprocess(image)
+  unsqueezed_image = processed_image.unsqueeze(0).to(device)
+  embed_image = model.encode_image(unsqueezed_image)
+
+  # Detach, move to CPU, convert to numpy array, and extract the first element as a list
+  result = embed_image.detach().cpu().numpy()[0].tolist()
+  return result
+
+data = []
+for i in range(len(ds["train"])):
+    img = ds["train"][i]['image']
+    text = ds["train"][i]['text']
+
+    # Encode the image
+    encoded_img = embedding_image(img)
+    data.append({"vector": encoded_img, "text": text, "id" : i})
+
+
+db = lancedb.connect('./data/tables')
+schema = pa.schema(
+  [pa.field("vector", pa.list_(pa.float32(),768)),
+   pa.field("text", pa.string()),
+   pa.field("id", pa.int32()) ])
+
+tabel = db.create_table("GTA Image Embedding Data", schema=schema, mode="overwrite")
+
+tabel.add(data)
+tabel.to_pandas()
+
+import gradio as gr
+
+# Define your search_text function
+def search_text(text):
+    query = tabel.search(embedding_text(text)).limit(4).to_pandas()
+    images = []
+    for i in range(len(query)):
+        data_id = int(query['id'][i])
+        image_path = ds["train"][data_id]['image']
+        images.append(image_path)
+    return images
+
+# Create Gradio interface
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Tab("Image search"):
+            vector_query = gr.Textbox(value="Input Text to search; Car on traffic light ", show_label=False)
+            b1 = gr.Button("Submit")
+
+    with gr.Row():
+        gallery = gr.Gallery(
+            label="Found images", show_label=False, elem_id="gallery"
+        ).style(columns=[2], object_fit="contain", height="auto")
+
+    b1.click(search_text, inputs=vector_query, outputs=gallery)
+
+demo.launch()
+
+
+
+
+

requirements.txt

@@ -0,0 +1,10 @@
+numpy
+pandas
+matplotlib
+gradio
+datasets
+torch
+clip
+PIL
+pyarrow
+lancedb