Update app.py

app.py

@@ -4,20 +4,21 @@ import pandas as pd
 import torch
 import gradio as gr
 
+# Import dataset loader
+from datasets import load_dataset
+
 # Import CLIP model and processor
 from transformers import CLIPModel, CLIPProcessor
 
-# Import image handling
-from PIL import Image
-
 
 # Select device (GPU if available)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 # Load pretrained CLIP model
-model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+MODEL_NAME = "openai/clip-vit-base-patch32"
+model = CLIPModel.from_pretrained(MODEL_NAME)
+processor = CLIPProcessor.from_pretrained(MODEL_NAME)
 
 # Move model to device and set evaluation mode
 model = model.to(device)
@@ -27,24 +28,22 @@ model.eval()
 # Load precomputed embeddings from file
 emb_df = pd.read_parquet("clip_embeddings_3000.parquet")
 
-# Extract image identifiers
-sampled_ids = emb_df["image_id"].values
-
 # Extract normalized embeddings matrix
-embeddings = emb_df.drop(columns=["image_id"]).values.astype("float32")
+embeddings = emb_df.drop(columns=["image_id"]).values.astype(np.float32)
 
 
-# Attempt to load sampled indices (optional optimization)
-try:
-    sampled_indices = np.load("sampled_indices_3000.npy")
-except Exception:
-    sampled_indices = None
+# Load sampled indices (required to fetch the same 3000 images)
+sampled_indices = np.load("sampled_indices_3000.npy").astype(int).tolist()
+
+# Load dataset and select the sampled subset
+ds = load_dataset("JamieSJS/stanford-online-products", "corpus")["corpus"]
+sampled_dataset = ds.select(sampled_indices)
 
 
 # Convert a user image into a normalized CLIP embedding
-def embed_image(image: Image.Image):
+def embed_image(image):
     # Preprocess image for CLIP
-    inputs = processor(images=image, return_tensors="pt")
+    inputs = processor(images=[image], return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
     # Extract image features without gradients
@@ -52,42 +51,51 @@ def embed_image(image: Image.Image):
         features = model.get_image_features(**inputs)
 
     # Convert embedding to numpy and normalize
-    vec = features.cpu().numpy()[0]
-    vec = vec / np.linalg.norm(vec)
+    vec = features.cpu().numpy().reshape(-1).astype(np.float32)
+    vec = vec / (np.linalg.norm(vec) + 1e-12)
 
     return vec
 
 
 # Recommend top-3 visually similar images
 def recommend(image):
-    # Embed user input image
-    user_vec = embed_image(image)
+    try:
+        # Embed user input image
+        user_vec = embed_image(image)
+
+        # Compute cosine similarity scores
+        scores = embeddings @ user_vec
 
-    # Compute cosine similarity scores
-    scores = embeddings @ user_vec
+        # Get Top-3 indices
+        top_idx = np.argsort(scores)[::-1][:3]
+        top_scores = scores[top_idx]
 
-    # Retrieve top-4 results (including query itself)
-    top_idx = np.argsort(scores)[::-1][:4]
+        # Fetch images directly from the sampled dataset
+        results = [sampled_dataset[int(i)]["image"] for i in top_idx]
 
-    # Remove the first result (query image)
-    top_idx = top_idx[1:]
+        # Optional: return a short message for visibility
+        msg = (
+            f"Top-3 cosine similarity scores: "
+            f"{top_scores[0]:.3f}, {top_scores[1]:.3f}, {top_scores[2]:.3f}"
+        )
 
-    # Load recommended images
-    results = []
-    for idx in top_idx:
-        img = Image.open(sampled_ids[idx]).convert("RGB")
-        results.append(img)
+        return results, msg
 
-    return results
+    except Exception as e:
+        return [], f"Error: {str(e)}"
 
 
 # Define Gradio interface
 demo = gr.Interface(
     fn=recommend,
     inputs=gr.Image(type="pil", label="Upload an image"),
-    outputs=gr.Gallery(label="Top-3 Recommended Images"),
+    outputs=[
+        gr.Gallery(label="Top-3 Recommended Images"),
+        gr.Textbox(label="Details"),
+    ],
     title="CLIP Image Recommendation System",
-    description="Upload an image and receive visually similar product recommendations."
+    description="Upload an image and receive visually similar product recommendations.",
+    allow_flagging="never",
 )