Update app.py

app.py

@@ -2,8 +2,6 @@ import streamlit as st
 from datasets import load_dataset
 from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
-from io import BytesIO
-import requests
 import random
 
 # Load the CLIP model and processor
@@ -25,21 +23,25 @@ def load_streamed_dataset():
 dataset = load_streamed_dataset()
 
 def fetch_random_memes():
-    sample_size = 100
+    """Fetch two random memes from the dataset."""
+    sample_size = 100  # Number of samples to shuffle
     dataset_samples = list(dataset.shuffle(seed=random.randint(0, 1000)).take(sample_size))
     meme1, meme2 = random.sample(dataset_samples, 2)
     return meme1, meme2
 
 def parse_meme(meme):
+    """Extract the caption and Pillow image from a meme."""
     caption = meme["answers"][0] if meme.get("answers") else "No caption available"
-    image_url = meme["image"]
-    return caption, image_url
+    image = meme["image"]  # This is already a PIL image object
+    return caption, image
 
-def score_meme(image_url, caption):
+def score_meme(image, caption):
+    """Score a meme by evaluating the image-caption compatibility."""
     try:
-        response = requests.get(image_url)
-        image = Image.open(BytesIO(response.content)).convert("RGB")
+        # Preprocess image and caption
         inputs = processor(text=[caption], images=[image], return_tensors="pt", padding=True)
+
+        # Get the compatibility score
         outputs = model(**inputs)
         logits_per_text = outputs.logits_per_text
         return logits_per_text.item()
@@ -48,25 +50,31 @@ def score_meme(image_url, caption):
         return 0
 
 if st.button("Start Meme Battle"):
+    # Fetch random memes
     meme1, meme2 = fetch_random_memes()
-    caption1, image_url1 = parse_meme(meme1)
-    caption2, image_url2 = parse_meme(meme2)
 
-    score1 = score_meme(image_url1, caption1)
-    score2 = score_meme(image_url2, caption2)
+    # Parse captions and images
+    caption1, image1 = parse_meme(meme1)
+    caption2, image2 = parse_meme(meme2)
+
+    # Score memes
+    score1 = score_meme(image1, caption1)
+    score2 = score_meme(image2, caption2)
 
+    # Display Meme 1
     st.write("#### Meme 1")
-    st.image(image_url1, caption=f"Caption: {caption1}")
+    st.image(image1, caption=f"Caption: {caption1}")
     st.write(f"AI Score: {score1:.2f}")
 
+    # Display Meme 2
     st.write("#### Meme 2")
-    st.image(image_url2, caption=f"Caption: {caption2}")
+    st.image(image2, caption=f"Caption: {caption2}")
     st.write(f"AI Score: {score2:.2f}")
 
+    # Determine the winner
     if score1 > score2:
         st.write("🎉 **Meme 1 Wins!**")
     elif score2 > score1:
         st.write("🎉 **Meme 2 Wins!**")
     else:
         st.write("🤝 **It's a tie!**")
-