Update app.py

app.py

@@ -1,4 +1,4 @@
-# 1. Install Required Libraries (if you haven't already)
+# 1. Install Required Libraries (run once in terminal or notebook)
 # pip install gradio transformers torch torchvision pillow
 
 # 2. Import Libraries
@@ -7,52 +7,52 @@ from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
 import torch
 
-# 3. Load the CLIP model
+# 3. Load the Pre-trained CLIP Model
 model_name = "openai/clip-vit-base-patch16"
 processor = CLIPProcessor.from_pretrained(model_name)
 model = CLIPModel.from_pretrained(model_name)
 
-# 4. Define the matching function
+# 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
+    # Parse user input into a list of captions
     captions = [line.strip() for line in descriptions.strip().split('\n') if line.strip()]
     
     if len(captions) < 2:
         return "Please enter at least two descriptions to compare."
 
-    # Process inputs
+    # Preprocess inputs
     inputs = processor(text=captions, images=image, return_tensors="pt", padding=True)
     
-    # Run model
+    # Run CLIP model
     with torch.no_grad():
         outputs = model(**inputs)
 
-    # Calculate probabilities
+    # Get prediction scores
     logits_per_image = outputs.logits_per_image  # shape: [1, num_captions]
-    probs = logits_per_image.softmax(dim=1)[0]    # Convert to probabilities
+    probs = logits_per_image.softmax(dim=1)[0]   # shape: [num_captions]
 
-    # Build result dictionary
-    results = {captions[i]: f"{probs[i].item() * 100:.2f}%" for i in range(len(captions))}
-    
-    # Sort by confidence (descending)
-    sorted_results = dict(sorted(results.items(), key=lambda item: float(item[1][:-1]), reverse=True))
-    
-    return sorted_results
+    # Format results into a dictionary with raw floats (0.0 - 1.0)
+    result_dict = {captions[i]: probs[i].item() for i in range(len(captions))}
+
+    # Pick the best match
+    best_caption = max(result_dict, key=result_dict.get)
+
+    return best_caption, result_dict
 
-# 5. Create the Gradio interface
+# 5. Build Gradio Interface
 iface = gr.Interface(
     fn=match_image_with_descriptions,
     inputs=[
         gr.Image(type="pil", label="Upload an Image"),
         gr.Textbox(lines=6, placeholder="Enter one description per line...", label="Descriptions")
     ],
-    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."
+    outputs=gr.Label(label="Best Match with Confidence Scores"),
+    title="🧠 CLIP Image-Text Matcher",
+    description="Upload an image and enter multiple captions (one per line). The AI will compare them and show which caption best fits the image.",
 )
 
-# 6. Launch the app
+# 6. Launch the App
 iface.launch()