Update app.py

app.py

@@ -6,18 +6,19 @@ import faiss
 import numpy as np
 from huggingface_hub import hf_hub_download
 from datasets import load_dataset
+import pandas as pd
 
-hf_hub_download("merve/siglip-faiss-wikiart", "siglip_new.index", local_dir="./")
-index = faiss.read_index("./siglip_new.index")
+# download model and dataset
+hf_hub_download("merve/siglip-faiss-wikiart", "siglip_10k.index", local_dir="./")
+hf_hub_download("merve/siglip-faiss-wikiart", "wikiart_10k.csv", local_dir="./")
 
-dataset = load_dataset("huggan/wikiart")
+# read index, dataset and load siglip model and processor
+index = faiss.read_index("./siglip_10k.index")
+df = pd.read_csv("./wikiart_10k.csv")
 device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
-dataset = dataset.with_format("torch", device=device)
-
 processor = AutoProcessor.from_pretrained("nielsr/siglip-base-patch16-224")
 model = SiglipModel.from_pretrained("nielsr/siglip-base-patch16-224").to(device)
 
-
 def extract_features_siglip(image):
     with torch.no_grad():
         inputs = processor(images=image, return_tensors="pt").to(device)
@@ -29,14 +30,17 @@ def infer(input_image):
   input_features = input_features.detach().cpu().numpy()
   input_features = np.float32(input_features)
   faiss.normalize_L2(input_features)
-  distances, indices = index2.search(input_features, 9)
+  distances, indices = index.search(input_features, 3)
   gallery_output = []
   for i,v in enumerate(indices[0]):
     sim = -distances[0][i]
-    img_resized = dataset["train"][int(v)]['image']
-    gallery_output.append(img_resized)
+    image_url = df.iloc[v]["Link"]
+    img_retrieved = read_image_from_url(image_url)
+    gallery_output.append(img_retrieved)
   return gallery_output
 
+
+description="This is an application where you can draw an image and find the closest artwork among 10k art from wikiart dataset."
 gr.Interface(infer, "sketchpad", "gallery",  title="Draw to Search Art 🖼️").launch()