import os import numpy as np import pandas as pd import torch import gradio as gr import imagehash from PIL import Image from datasets import load_dataset from transformers import CLIPModel, CLIPProcessor # ----------------------------- # Configuration # ----------------------------- DATASET_NAME = "Leonardo6/memecap" CLIP_MODEL_NAME = "openai/clip-vit-base-patch32" FINAL_IMAGE_RECOMMENDER_SETTINGS = { "candidate_pool_size": 150, "clip_weight": 0.60, "hash_weight": 0.35, "layout_weight": 0.05, "strong_hash_threshold": 0.78 } TOP_K = 3 # ----------------------------- # Device # ----------------------------- device = "cuda" if torch.cuda.is_available() else "cpu" # ----------------------------- # Load saved deployment files # ----------------------------- image_embeddings_normalized = np.load("image_embeddings.npy") image_indices = np.load("image_indices.npy") meme_metadata = pd.read_csv("meme_metadata.csv") layout_df = pd.read_csv("layout_features.csv") def load_hash_file(file_path): with open(file_path, "r") as f: return [imagehash.hex_to_hash(line.strip()) for line in f.readlines()] image_phashes = load_hash_file("image_phashes.txt") image_dhashes = load_hash_file("image_dhashes.txt") image_whashes = load_hash_file("image_whashes.txt") # ----------------------------- # Load dataset and model # ----------------------------- dataset = load_dataset(DATASET_NAME) memes = dataset["train"] clip_model = CLIPModel.from_pretrained(CLIP_MODEL_NAME).to(device) clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME) clip_model.eval() # ----------------------------- # Helper functions # ----------------------------- def get_dataset_image(dataset_index): image = memes[int(dataset_index)]["images"][0] if not isinstance(image, Image.Image): image = Image.fromarray(image) return image.convert("RGB") def normalize_vector(vector): vector = np.asarray(vector) norm = np.linalg.norm(vector) if norm == 0: return vector return vector / norm def extract_tensor_from_model_output(output): if isinstance(output, torch.Tensor): return output if hasattr(output, "pooler_output") and output.pooler_output is not None: return output.pooler_output if hasattr(output, "last_hidden_state") and output.last_hidden_state is not None: return output.last_hidden_state[:, 0, :] raise ValueError("Could not extract tensor from model output.") def hash_similarity_from_hash_lists(query_hashes, candidate_indices): similarities = [] query_phash, query_dhash, query_whash = query_hashes for candidate_idx in candidate_indices: candidate_idx = int(candidate_idx) phash_sim = 1 - ((query_phash - image_phashes[candidate_idx]) / 64) dhash_sim = 1 - ((query_dhash - image_dhashes[candidate_idx]) / 64) whash_sim = 1 - ((query_whash - image_whashes[candidate_idx]) / 64) avg_sim = (phash_sim + dhash_sim + whash_sim) / 3 similarities.append(avg_sim) return np.array(similarities) def layout_similarity_from_uploaded(uploaded_layout, candidate_indices): query_aspect_ratio = uploaded_layout["aspect_ratio"] candidate_layouts = layout_df.iloc[candidate_indices] candidate_aspect_ratios = candidate_layouts["aspect_ratio"].values aspect_diff = np.abs(candidate_aspect_ratios - query_aspect_ratio) similarities = 1 / (1 + aspect_diff) return similarities def get_uploaded_image_embedding(uploaded_image): if uploaded_image is None: return None uploaded_image = uploaded_image.convert("RGB") inputs = clip_processor( images=uploaded_image, return_tensors="pt" ).to(device) with torch.no_grad(): output = clip_model.get_image_features(**inputs) image_features = extract_tensor_from_model_output(output) image_embedding = image_features.detach().cpu().numpy()[0] image_embedding = normalize_vector(image_embedding) return image_embedding def get_uploaded_image_hashes(uploaded_image): uploaded_image = uploaded_image.convert("RGB") uploaded_phash = imagehash.phash(uploaded_image) uploaded_dhash = imagehash.dhash(uploaded_image) uploaded_whash = imagehash.whash(uploaded_image) return uploaded_phash, uploaded_dhash, uploaded_whash def get_uploaded_image_layout(uploaded_image): uploaded_image = uploaded_image.convert("RGB") width, height = uploaded_image.size aspect_ratio = width / height if height != 0 else 0 return { "width": width, "height": height, "aspect_ratio": aspect_ratio } # ----------------------------- # Image upload recommendation # ----------------------------- def recommend_by_uploaded_image(uploaded_image, top_k=TOP_K): if uploaded_image is None: return [] settings = FINAL_IMAGE_RECOMMENDER_SETTINGS query_embedding = get_uploaded_image_embedding(uploaded_image) uploaded_hashes = get_uploaded_image_hashes(uploaded_image) uploaded_layout = get_uploaded_image_layout(uploaded_image) clip_similarities = image_embeddings_normalized @ query_embedding candidate_pool_size = min(settings["candidate_pool_size"], len(clip_similarities)) candidate_indices = np.argsort(clip_similarities)[::-1][:candidate_pool_size] candidate_clip_scores = clip_similarities[candidate_indices] candidate_hash_scores = hash_similarity_from_hash_lists(uploaded_hashes, candidate_indices) candidate_layout_scores = layout_similarity_from_uploaded(uploaded_layout, candidate_indices) final_scores = ( settings["clip_weight"] * candidate_clip_scores + settings["hash_weight"] * candidate_hash_scores + settings["layout_weight"] * candidate_layout_scores ) results = pd.DataFrame({ "dataset_index": candidate_indices, "clip_score": candidate_clip_scores, "hash_score": candidate_hash_scores, "layout_score": candidate_layout_scores, "final_score": final_scores }) results["match_type"] = np.where( results["hash_score"] >= settings["strong_hash_threshold"], "strong_template_match", "fallback_best_available" ) results = results.sort_values( by=["match_type", "final_score"], ascending=[False, False] ) selected_results = results.head(top_k).copy() selected_results["rank"] = range(1, len(selected_results) + 1) return selected_results # ----------------------------- # Text query recommendation # ----------------------------- def get_text_query_embedding(text_query): inputs = clip_processor( text=[text_query], return_tensors="pt", padding=True, truncation=True ).to(device) with torch.no_grad(): output = clip_model.get_text_features(**inputs) text_features = extract_tensor_from_model_output(output) text_embedding = text_features.detach().cpu().numpy()[0] text_embedding = normalize_vector(text_embedding) return text_embedding def recommend_by_text_query(text_query, top_k=TOP_K): if text_query is None or text_query.strip() == "": return [] text_embedding = get_text_query_embedding(text_query) similarities = image_embeddings_normalized @ text_embedding top_indices = np.argsort(similarities)[::-1][:top_k] results = pd.DataFrame({ "dataset_index": top_indices, "similarity_score": similarities[top_indices], "match_type": "text_match" }) results["rank"] = range(1, len(results) + 1) return results # ----------------------------- # Gradio formatting # ----------------------------- def format_results_for_gradio(results): if results is None or len(results) == 0: return [] gallery_items = [] for _, row in results.iterrows(): dataset_index = int(row["dataset_index"]) image = get_dataset_image(dataset_index) if "final_score" in row: label = ( f"Rank {int(row['rank'])} | " f"Index {dataset_index} | " f"Final score: {row['final_score']:.3f} | " f"{row['match_type']}" ) else: label = ( f"Rank {int(row['rank'])} | " f"Index {dataset_index} | " f"Similarity score: {row['similarity_score']:.3f} | " f"{row['match_type']}" ) gallery_items.append((image, label)) return gallery_items def gradio_recommend_by_uploaded_image(uploaded_image): results = recommend_by_uploaded_image(uploaded_image, top_k=TOP_K) return format_results_for_gradio(results) def gradio_recommend_by_text_query(text_query): results = recommend_by_text_query(text_query, top_k=TOP_K) return format_results_for_gradio(results) # ----------------------------- # Gradio app # ----------------------------- description = """ # MemeMatch: Visual and Template-Based Meme Recommendation System MemeMatch is a recommendation system that helps users find meme templates using either an uploaded image or a text description. The system uses CLIP embeddings for visual-semantic similarity, perceptual hashes for template/layout support, and a small layout score based on aspect ratio. ### How to use **Image Upload:** upload a meme image to find visually similar meme templates. **Text Search:** describe the kind of meme template you want using a visual prompt. The app returns the top 3 recommendations with similarity scores and match labels. """ method_note = """ ### Match Labels - **strong_template_match**: the result has strong visual/template similarity based on perceptual hash support. - **fallback_best_available**: the system did not find enough strong template matches, so it returns the best available results from the dataset. This is intentional. Some meme templates appear only once or twice in the dataset, so not every query can produce three strong matches. """ text_examples = """ ### Strong text prompts Try prompts like: - `a meme template showing a character giving a presentation in front of an audience` - `a meme showing people reacting in an audience` - `a meme template with Spongebob characters` - `a meme showing someone pressing a button` Text search works best when the prompt describes what the image should visually contain. """ limitations_note = """ ### Project Limitation Image-to-image recommendation is not perfect because CLIP captures broad visual meaning, while perceptual hashes capture low-level image structure. The hybrid method balances both, but weak fallback results can still appear when the dataset does not contain enough similar meme templates. """ with gr.Blocks(title="MemeMatch") as demo: gr.Markdown(description) with gr.Accordion("Method Notes and Limitations", open=False): gr.Markdown(method_note) gr.Markdown(limitations_note) with gr.Tab("Image Upload"): gr.Markdown( """ ## Image Upload Recommendation Upload a meme image and click the button below. The system compares your image to the dataset using a hybrid score: `CLIP similarity + perceptual hash similarity + weak layout similarity` """ ) image_input = gr.Image( type="pil", label="Upload a meme image" ) image_button = gr.Button("Find Similar Meme Templates") image_output = gr.Gallery( label="Top 3 Image Recommendations", columns=3, rows=1, height=350 ) image_button.click( fn=gradio_recommend_by_uploaded_image, inputs=image_input, outputs=image_output ) with gr.Tab("Text Search"): gr.Markdown( """ ## Text-to-Image Meme Search Describe the meme template you want. Use a visual description, not a vague category. """ ) gr.Markdown(text_examples) text_input = gr.Textbox( label="Describe the meme template you want", placeholder="Example: a meme showing someone pressing a button", lines=2 ) text_button = gr.Button("Search Meme Templates") text_output = gr.Gallery( label="Top 3 Text-Based Recommendations", columns=3, rows=1, height=350 ) text_button.click( fn=gradio_recommend_by_text_query, inputs=text_input, outputs=text_output ) if __name__ == "__main__": demo.launch()