Update app.py

app.py

@@ -4,13 +4,12 @@ from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
 import os
 import faiss
-from datasets import load_dataset, concatenate_datasets
-import requests
+from datasets import load_dataset
+import requests 
 import io
-import time
 
 # --- Configuration ---
-MODEL_PATH = "clip_finetuned"
+MODEL_PATH = "clip_finetuned" 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 FAISS_INDEX_PATH = "gallery.index"
 
@@ -22,79 +21,47 @@ processor = CLIPProcessor.from_pretrained(MODEL_PATH)
 print("Loading FAISS index...")
 faiss_index = faiss.read_index(FAISS_INDEX_PATH)
 
-# --- Load Full COCO Training and Validation Datasets ---
-print("Connecting to COCO dataset (train and validation splits) on the Hub...")
-train_dataset = load_dataset("phiyodr/coco2017", split="train")
-val_dataset = load_dataset("phiyodr/coco2017", split="validation")
-combined_dataset = concatenate_datasets([train_dataset, val_dataset])
-print(f"Successfully connected to combined dataset with {len(combined_dataset)} images.")
-
-# --- Filter for Child-Friendly Content ---
-def filter_child_friendly(dataset):
-    adult_keywords = ["nude", "violence", "adult", "gun", "blood"]
-    filtered_dataset = []
-    for item in dataset:
-        file_name = item.get('file_name', '').lower()
-        # Exclude images with adult-related keywords in file_name
-        if not any(keyword in file_name for keyword in adult_keywords):
-            filtered_dataset.append(item)
-    return filtered_dataset
-
-filtered_dataset = filter_child_friendly(combined_dataset)
-print(f"Filtered dataset size for child-friendly content: {len(filtered_dataset)} images.")
+# --- Connect to the COCO dataset on the Hub ---
+print("Connecting to COCO dataset on the Hub...")
+val_dataset = load_dataset("phiyodr/coco2017", split="validation", trust_remote_code=True)
+    
+print(f"Successfully connected to dataset with {len(val_dataset)} images.")
 
-# --- The Search Function with Metrics ---
+# --- The Search Function (Corrected) ---
 def image_search(query_text: str, top_k: int):
-    start_time = time.time()
     with torch.no_grad():
         inputs = processor(text=query_text, return_tensors="pt").to(DEVICE)
         text_embedding = model.get_text_features(**inputs)
         text_embedding /= text_embedding.norm(p=2, dim=-1, keepdim=True)
 
     distances, indices = faiss_index.search(text_embedding.cpu().numpy(), int(top_k))
-
-    # Process results with metrics
     results = []
-    relevant_count = 0
-    retrieval_time = time.time() - start_time
-    memory_usage = torch.cuda.memory_allocated() / 1024**2 if DEVICE == "cuda" else os.cpu_count() * 10  # Approx. MB
-
     for i in indices[0]:
-        if i < len(filtered_dataset):
-            item = filtered_dataset[int(i)]
-            image_url = item['coco_url']  # Assuming coco_url is available
-            response = requests.get(image_url)
-            image = Image.open(io.BytesIO(response.content)).convert("RGB")
-            results.append(image)
-            # Simple relevance check based on file_name matching query
-            file_name = item.get('file_name', '').lower()
-            if query_text.lower() in file_name:
-                relevant_count += 1
-
-    accuracy = (relevant_count / top_k) * 100 if top_k > 0 else 0
-    metrics = f"Retrieval Time: {retrieval_time:.2f} seconds, Memory Usage: {memory_usage:.2f} MB, Accuracy: {accuracy:.2f}%"
-
-    print(metrics)
-    return results, metrics
-
-# --- Gradio Interface ---
+        item = val_dataset[int(i)]
+        image_url = item['coco_url']
+        response = requests.get(image_url)
+        image = Image.open(io.BytesIO(response.content)).convert("RGB")
+        results.append(image)
+        
+    return results
+
+# --- Gradio Interface (No changes needed here) ---
 with gr.Blocks(theme=gr.themes.Soft()) as iface:
     gr.Markdown("# 🖼️ CLIP-Powered Image Search Engine")
-    gr.Markdown("Enter a text description to search for child-friendly images from the COCO dataset.")
+    gr.Markdown("Enter a text description to search for matching images.")
     
     with gr.Row():
-        query_input = gr.Textbox(label="Search Query", placeholder="e.g., a dog playing", scale=4)
+        query_input = gr.Textbox(label="Search Query", placeholder="e.g., a red car parked near a building", scale=4)
         k_slider = gr.Slider(minimum=1, maximum=12, value=4, step=1, label="Number of Results")
         submit_btn = gr.Button("Search", variant="primary")
 
     gallery_output = gr.Gallery(label="Search Results", show_label=False, columns=4, height="auto")
-    metrics_output = gr.Textbox(label="Performance Metrics", interactive=False)
 
-    submit_btn.click(fn=image_search, inputs=[query_input, k_slider], outputs=[gallery_output, metrics_output])
+    submit_btn.click(fn=image_search, inputs=[query_input, k_slider], outputs=gallery_output)
     
     gr.Examples(
-        examples=[["a dog playing", 4], ["children in a park", 8]],
+        examples=[["a dog catching a frisbee", 4], ["two people eating pizza", 8]],
         inputs=[query_input, k_slider]
     )
 
-iface.launch(share=True)
+iface.launch()