Update app.py

app.py

@@ -26,10 +26,10 @@ processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 # Load dataset
 dataset = load_dataset("lirus18/deepfashion", split="train")
 
-# Embed a subset of images
+# Embed a subset of dataset images
 image_vectors = []
 image_indices = []
-N = 500
+N = 500  # You can increase this if performance is okay
 
 for i in range(N):
     img = dataset[i]['image'].convert("RGB")
@@ -41,7 +41,7 @@ for i in range(N):
 
 image_vectors = np.array(image_vectors)
 
-# Find similar images
+# Similarity search
 def find_similar(user_image, top_k=3, exclude_index=None):
     inputs = processor(images=user_image.convert("RGB"), return_tensors="pt").to(device)
     with torch.no_grad():
@@ -52,7 +52,7 @@ def find_similar(user_image, top_k=3, exclude_index=None):
         sims[exclude_index] = -1
 
     top_idx = sims.argsort()[-top_k:][::-1]
-    return [dataset[image_indices[i]]['image'] for i in top_idx]
+    return [dataset[image_indices[i]]['image'] for i in top_idx], query_vec
 
 # Load Stable Diffusion
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
@@ -73,8 +73,11 @@ def generate_outfits(input_image, n=10):
 
     return generated_images
 
+# Main function
 def recommend_from_upload(uploaded_image):
     uploaded_image = uploaded_image.convert("RGB")
+
+    # Check if the uploaded image already exists in the dataset
     closest_idx = None
     for i in range(len(image_indices)):
         dataset_image = dataset[image_indices[i]]['image'].convert("RGB")
@@ -82,12 +85,28 @@ def recommend_from_upload(uploaded_image):
             closest_idx = i
             break
 
-    similar_imgs = find_similar(uploaded_image, top_k=3, exclude_index=closest_idx)
-    generated_imgs = generate_outfits(uploaded_image, n=3)
+    # Find 3 similar dataset items + get uploaded image embedding
+    similar_imgs, query_vec = find_similar(uploaded_image, top_k=3, exclude_index=closest_idx)
+
+    # Generate 10 new outfit images
+    generated_imgs = generate_outfits(uploaded_image, n=10)
 
-    return [uploaded_image] + similar_imgs + generated_imgs
+    # Find the most relevant generated image
+    best_score = -1
+    best_generated_img = None
+    for img in generated_imgs:
+        inputs = processor(images=img, return_tensors="pt").to(device)
+        with torch.no_grad():
+            emb = model.get_image_features(**inputs).cpu().numpy()
+        sim = cosine_similarity(query_vec, emb)[0][0]
+        if sim > best_score:
+            best_score = sim
+            best_generated_img = img
 
-# 5 clickable example images (must be uploaded to the repo)
+    # Final output: input + 3 similar + 1 most relevant generated
+    return [uploaded_image] + similar_imgs + [best_generated_img]
+
+# Example images (make sure they are in the root folder)
 example_paths = [
     ["example1.jpg"],
     ["example2.jpg"],
@@ -96,8 +115,6 @@ example_paths = [
     ["example5.jpg"]
 ]
 
-
-
 # Gradio Interface
 demo = gr.Interface(
     fn=recommend_from_upload,
@@ -107,14 +124,14 @@ demo = gr.Interface(
         gr.Image(label="Similar Item 1"),
         gr.Image(label="Similar Item 2"),
         gr.Image(label="Similar Item 3"),
-        gr.Gallery(label="AI-Generated Outfits (x10)", columns=5, rows=2),
-
+        gr.Image(label="Best AI-Generated Outfit")
     ],
     title="👗 Fashion Outfit Recommender",
-    description="Upload a clothing image to get 3 similar items from the dataset and 10 AI-generated outfit designs.",
+    description="Upload a clothing image to get 3 similar items from the dataset and 1 best AI-generated outfit design.",
     examples=example_paths
 )
 
 if __name__ == "__main__":
     demo.launch()
 
+