neojack/memebattle

1import streamlit as st2from datasets import load_dataset3from transformers import CLIPProcessor, CLIPModel4from PIL import Image5import random6 7# Load the CLIP model and processor8st.title("Meme Battle AI")9st.write("Stream memes directly and let AI determine the winner!")10 11@st.cache_resource12def load_model():13    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")14    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")15    return model, processor16 17model, processor = load_model()18 19@st.cache_resource20def load_streamed_dataset():21    return load_dataset("Dhruv-goyal/memes_with_captions", split="train", streaming=True)22 23dataset = load_streamed_dataset()24 25def fetch_random_memes():26    """Fetch two random memes from the dataset."""27    sample_size = 100  # Number of samples to shuffle28    dataset_samples = list(dataset.shuffle(seed=random.randint(0, 1000)).take(sample_size))29    meme1, meme2 = random.sample(dataset_samples, 2)30    return meme1, meme231 32def parse_meme(meme):33    """Extract the caption and Pillow image from a meme."""34    caption = meme["answers"][0] if meme.get("answers") else "No caption available"35    image = meme["image"]  # This is already a PIL image object36    return caption, image37 38def score_meme(image, caption):39    """Score a meme by evaluating the image-caption compatibility."""40    try:41        # Preprocess image and caption42        inputs = processor(text=[caption], images=[image], return_tensors="pt", padding=True)43 44        # Get the compatibility score45        outputs = model(**inputs)46        logits_per_text = outputs.logits_per_text47        return logits_per_text.item()48    except Exception as e:49        st.error(f"Error scoring meme: {e}")50        return 051 52if st.button("Start Meme Battle"):53    # Fetch random memes54    meme1, meme2 = fetch_random_memes()55 56    # Parse captions and images57    caption1, image1 = parse_meme(meme1)58    caption2, image2 = parse_meme(meme2)59 60    # Score memes61    score1 = score_meme(image1, caption1)62    score2 = score_meme(image2, caption2)63 64    # Display Meme 1 and Meme 2 side by side65    col1, col2 = st.columns(2)66 67    with col1:68        st.write("#### Meme 1")69        st.image(image1, caption=f"Caption: {caption1}")70        st.write(f"AI Score: {score1:.2f}")71 72    with col2:73        st.write("#### Meme 2")74        st.image(image2, caption=f"Caption: {caption2}")75        st.write(f"AI Score: {score2:.2f}")76 77    # Determine the winner78    if score1 > score2:79        st.write("🎉 **Meme 1 Wins!**")80    elif score2 > score1:81        st.write("🎉 **Meme 2 Wins!**")82    else:83        st.write("🤝 **It's a tie!**")84 85