app.py

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()