Update app.py

app.py

@@ -1,4 +1,4 @@
-# 1. Install Required Libraries (run this in terminal or notebook once)
+# 1. Install Required Libraries (if you haven't already)
 # pip install gradio transformers torch torchvision pillow
 
 # 2. Import Libraries
@@ -7,44 +7,52 @@ from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
 import torch
 
-# 3. Load the Pre-trained Model
+# 3. Load the CLIP model
 model_name = "openai/clip-vit-base-patch16"
 processor = CLIPProcessor.from_pretrained(model_name)
 model = CLIPModel.from_pretrained(model_name)
 
-# 4. Define the Prediction Function
-def classify_image_text(image, text):
-    if not image or not text:
-        return "Please provide both image and description."
+# 4. Define the matching function
+def match_image_with_descriptions(image, descriptions):
+    if not image or not descriptions.strip():
+        return "Please upload an image and enter at least one description."
+
+    # Convert multi-line string to list of captions
+    captions = [line.strip() for line in descriptions.strip().split('\n') if line.strip()]
     
-    # Process the inputs
-    inputs = processor(text=[text], images=image, return_tensors="pt", padding=True)
+    if len(captions) < 2:
+        return "Please enter at least two descriptions to compare."
+
+    # Process inputs
+    inputs = processor(text=captions, images=image, return_tensors="pt", padding=True)
     
-    # Get model predictions
+    # Run model
     with torch.no_grad():
         outputs = model(**inputs)
+
+    # Calculate probabilities
+    logits_per_image = outputs.logits_per_image  # shape: [1, num_captions]
+    probs = logits_per_image.softmax(dim=1)[0]    # Convert to probabilities
+
+    # Build result dictionary
+    results = {captions[i]: f"{probs[i].item() * 100:.2f}%" for i in range(len(captions))}
     
-    # Calculate similarity between image and text
-    logits_per_image = outputs.logits_per_image  # shape: [1, 1]
-    probs = logits_per_image.softmax(dim=1)      # shape: [1, 1]
-    score = probs[0][0].item()                   # Get scalar score
+    # Sort by confidence (descending)
+    sorted_results = dict(sorted(results.items(), key=lambda item: float(item[1][:-1]), reverse=True))
     
-    # Return readable percentage
-    match_percentage = round(score * 100, 2)
-    return f"Match Confidence: {match_percentage}%"
+    return sorted_results
 
-# 5. Create the Gradio Interface
+# 5. Create the Gradio interface
 iface = gr.Interface(
-    fn=classify_image_text,
+    fn=match_image_with_descriptions,
     inputs=[
         gr.Image(type="pil", label="Upload an Image"),
-        gr.Textbox(lines=2, placeholder="Describe the image...", label="Your Description")
+        gr.Textbox(lines=6, placeholder="Enter one description per line...", label="Descriptions")
     ],
-    outputs=gr.Label(label="Result"),
-    title="CLIP Image-Text Matcher",
-    description="Upload an image and enter a description. This app will tell you how well your text matches the image.",
-    allow_flagging="never"
+    outputs=gr.Label(label="Match Confidence"),
+    title="AI Image-Text Matcher",
+    description="Upload an image and enter multiple possible descriptions (one per line). The AI will tell you which one best matches the image."
 )
 
-# 6. Launch the App
+# 6. Launch the app
 iface.launch()