Hugging Face's logo

Spaces:
Unspoiled-Egg
/
Meme-Recommender
Sleeping

App Files Community
Meme-Recommender
File size: 3,706 Bytes
162175a
 
052f26d
 
162175a
052f26d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
162175a
 
052f26d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40fca10
dfad1ba
40fca10
 
052f26d
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import os
import numpy as np
import gradio as gr
from PIL import Image
from pathlib import Path

from config import settings
from utils.model_loading_util import load_model_from_hf
from src.processing.embedder import EmbeddingGenerator
from src.search.pinecone_indexer import PineconeImageSimilaritySearch


print("Initializing model and search index...")
model = load_model_from_hf(
    repo_id=settings.model_repo_id,
    device=settings.device,
    token=settings.hf_token
)


generator = EmbeddingGenerator(
    model=model,
    device=settings.device,
    batch_size=settings.batch_size,
    num_workers=2,
    layer_strategy="last_four_concat"
)

searcher = PineconeImageSimilaritySearch(
    index_name=settings.pinecone_index_name,
    api_key=settings.pinecone_api_key,
    dimension=settings.embedding_dim,
    metric="cosine",
)


def search_similar_images(input_img):
    if input_img is None:
        return None

    # Convert Gradio input (can be numpy array) to PIL Image
    if isinstance(input_img, np.ndarray):
        input_img = Image.fromarray(input_img)

    query_embedding = generator.generate_single_embedding(input_img)
    results = searcher.search(query_embedding, k=5, return_scores=True)

    gallery_items = []
    for path, score in results:
        img_path = os.path.join(settings.mount_path, path.lstrip("/"))
        img_path = Path(img_path)

        if img_path.exists():
            gallery_items.append((str(img_path), f"Similarity: {score:.4f}"))
        else:
            print(f"Warning: Image path not found: {img_path}")

    return gallery_items


custom_css = """
.container {
    max-width: 1000px;
    margin: auto;
    padding: 20px;
}
.header {
    text-align: center;
    margin-bottom: 30px;
}
.header h1 {
    font-size: 2.5rem;
    font-weight: 800;
    background: linear-gradient(90deg, #4F46E5, #EC4899);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 10px;
}
.header p {
    color: #6B7280;
    font-size: 1.1rem;
}
.gradio-container {
    background-color: #F9FAFB !important;
}
.gallery-container {
    border-radius: 12px;
    overflow: hidden;
    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
}
"""

with gr.Blocks() as demo:
    with gr.Column(elem_classes="container"):
        with gr.Column(elem_classes="header"):
            gr.Markdown("# Meme Similarity Search")
            gr.Markdown(
                "Upload an image to find the top 5 most similar memes in our database.")

        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(
                    label="Upload Image",
                    type="pil",
                    elem_id="input-img"
                )
                search_btn = gr.Button("Find Similar Memes", variant="primary")

            with gr.Column(scale=2):
                output_gallery = gr.Gallery(
                    label="Top 5 Similar Memes",
                    show_label=True,
                    elem_id="gallery",
                    columns=2,
                    rows=3,
                    object_fit="contain",
                    height="600px"
                )

        search_btn.click(
            fn=search_similar_images,
            inputs=input_image,
            outputs=output_gallery
        )

        input_image.upload(
            fn=search_similar_images,
            inputs=input_image,
            outputs=output_gallery
        )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        css=custom_css,
        allowed_paths=["/data"],
        show_error=True,
        debug=True
    )