Update app.py to add category and gender filtering

Implemented dropdowns for selecting clothing category and gender. Added filtering logic in the recommend function to return only results matching user selections. Also handled the case where no matches are found.

app.py

@@ -6,47 +6,76 @@ from PIL import Image
 from torchvision import transforms
 from huggingface_hub import hf_hub_download
 from sklearn.metrics.pairwise import cosine_similarity
+import timm
+
+# Load the precomputed fashion embeddings
+file_path = hf_hub_download(
+    repo_id="Elevi7/MatchMe",
+    filename="fashion_embeddings.pkl",
+    repo_type="dataset"
+)
 
-# Step 1: Load the precomputed fashion embeddings
-file_path = hf_hub_download(repo_id="Elevi7/MatchMe", filename="fashion_embeddings.pkl", repo_type="dataset")
 with open(file_path, "rb") as f:
     embedding_store = pickle.load(f)
 
-# Step 2: Load the same model used to create the embeddings
-# (update this if you used something different, e.g. CLIP, ResNet, etc.)
-import timm
+# Load the pretrained model (same one used for embeddings)
 model = timm.create_model("resnet18", pretrained=True)
 model.eval()
 
+# Define image transformation pipeline
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor()
 ])
 
+# Extract embedding from uploaded image
 def extract_embedding(image):
     image = transform(image).unsqueeze(0)
     with torch.no_grad():
         embedding = model(image).squeeze().numpy()
     return embedding
 
-def recommend(image):
+# Recommendation function with filtering
+def recommend(image, category, gender):
     query_embedding = extract_embedding(image)
 
-    # Get embeddings and paths
-    all_embeddings = np.array([item["embedding"] for item in embedding_store])
-    paths = [item["image_path"] for item in embedding_store]
+    # Filter embeddings by category and gender
+    filtered_store = [
+        item for item in embedding_store
+        if item.get("category") == category and item.get("gender") == gender
+    ]
+
+    if len(filtered_store) == 0:
+        print("No matches found for selected category and gender.")
+        return ["", "", ""]  # Return empty images
+
+    all_embeddings = np.array([item["embedding"] for item in filtered_store])
+    paths = [item["image_path"] for item in filtered_store]
 
     similarities = cosine_similarity([query_embedding], all_embeddings)[0]
-    top_indices = np.argsort(similarities)[-3:][::-1]  # Top 3 matches
+    top_indices = np.argsort(similarities)[-3:][::-1]  # Top 3
 
     return [paths[i] for i in top_indices]
 
+# Gradio interface
 demo = gr.Interface(
     fn=recommend,
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="filepath", label=f"Match {i+1}") for i in range(3)],
+    inputs=[
+        gr.Image(type="pil", label="Upload Clothing Image"),
+        gr.Dropdown(
+            choices=["shoes", "tops", "pants", "handbags", "coats_jackets",
+                     "sunglasses", "shorts", "skirts", "earrings", "necklaces"],
+            label="Category"
+        ),
+        gr.Dropdown(choices=["men", "women"], label="Gender")
+    ],
+    outputs=[
+        gr.Image(type="filepath", label="Match 1"),
+        gr.Image(type="filepath", label="Match 2"),
+        gr.Image(type="filepath", label="Match 3")
+    ],
     title="MatchMe: Fashion Recommender",
-    description="Upload a fashion image and get 3 visually similar items."
+    description="Upload a fashion image and get 3 visually similar items. Filter by category and gender."
 )
 
 demo.launch()