adding the files to run the app

app.py

@@ -0,0 +1,161 @@
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+from search import search_similarity, process_image_for_encoder_gradio
+from utils import str_to_bytes
+from io import BytesIO
+
+def add_ranking_number(image, rank):
+    """Añade un número de ranking a la imagen"""
+    img_with_rank = image.copy()
+    draw = ImageDraw.Draw(img_with_rank)
+    
+    width, height = image.size
+    circle_radius = min(width, height) // 15
+    circle_position = (circle_radius + 10, circle_radius + 10)
+    
+    draw.ellipse(
+        [(circle_position[0] - circle_radius, circle_position[1] - circle_radius),
+         (circle_position[0] + circle_radius, circle_position[1] + circle_radius)],
+        fill='white',
+        outline='black'
+    )
+    
+    font_size = circle_radius
+    try:
+        font = ImageFont.truetype("Arial.ttf", font_size)
+    except:
+        font = ImageFont.load_default()
+    
+    text = str(rank + 1)
+    text_bbox = draw.textbbox((0, 0), text, font=font)
+    text_width = text_bbox[2] - text_bbox[0]
+    text_height = text_bbox[3] - text_bbox[1]
+    text_position = (
+        circle_position[0] - text_width // 2,
+        circle_position[1] - text_height // 2
+    )
+    
+    draw.text(text_position, text, fill='black', font=font)
+    return img_with_rank
+
+def process_image_result(image_str, rank):
+    """Convierte una cadena de imagen en un objeto PIL Image con ranking"""
+    try:
+        img = Image.open(BytesIO(str_to_bytes(image_str)))
+        return add_ranking_number(img, rank)
+    except Exception as e:
+        print(f"Error procesando imagen: {e}")
+        return None
+      
+def interface_fn(mode, input_text, input_image, top_k):
+    try:
+        # Determinar qué input usar basado en el modo
+        if mode == "text":
+            if not input_text.strip():
+                return [], "Por favor, ingresa un texto para buscar."
+            input_data = input_text
+        else:  # mode == "image"
+            if input_image is None:
+                return [], "Por favor, sube una imagen para buscar."
+            input_data = process_image_for_encoder_gradio(input_image, is_bytes=False)
+
+        # Show the input data
+        print(f"Input data: {input_data}")  # Para debugging
+        
+        # Realizar la búsqueda
+        results = search_similarity(input_data, mode, int(top_k))
+        
+        # Formatear resultados según el modo
+        if mode == "text":  # Devuelve imágenes
+            processed_images = []
+            # Si results es una lista de listas, la aplanamos
+            if results and isinstance(results[0], list):
+                print("Recibida lista de listas, aplanando...")  # Para debugging
+                results = [item for sublist in results for item in sublist]
+            
+            for idx, img_str in enumerate(results):
+                img = process_image_result(img_str, idx)
+                if img is not None:
+                    processed_images.append(img)
+            
+            if not processed_images:
+                return [], "No se pudieron procesar las imágenes"
+            return processed_images, None
+            
+        else:  # mode == "image" - Devuelve textos
+            if isinstance(results, list):
+                numbered_texts = [f"{i+1}. {text}" for i, text in enumerate(results)]
+                return [], "\n\n".join(numbered_texts)
+            else:
+                return [], str(results)
+            
+    except Exception as e:
+        print(f"Error en interface_fn: {str(e)}")
+        print(f"Tipo de resultados: {type(results)}")  # Para debugging
+        return [], f"Error durante la búsqueda: {str(e)}"
+
+
+def search_text(input_text, top_k):
+    try:
+        if not input_text.strip():
+            return []
+        
+        # Realizar la búsqueda
+        results = search_similarity(input_text, "text", int(top_k))
+        
+        processed_images = []
+        # Si results es una lista de listas, la aplanamos
+        if results and isinstance(results[0], list):
+            results = [item for sublist in results for item in sublist]
+        
+        for idx, img_str in enumerate(results):
+            img = process_image_result(img_str, idx)
+            if img is not None:
+                processed_images.append(img)
+        
+        return processed_images
+            
+    except Exception as e:
+        print(f"Error en search_text: {str(e)}")
+        return []
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Buscador de Similitud por Texto")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_text = gr.Textbox(
+                label="Texto de búsqueda",
+                placeholder="Ingresa aquí tu texto...",
+                lines=3
+            )
+            
+            top_k = gr.Slider(
+                minimum=1,
+                maximum=20,
+                value=5,
+                step=1,
+                label="Número de resultados",
+                info="¿Cuántos resultados similares quieres ver?"
+            )
+            
+            search_button = gr.Button("Buscar")
+        
+        with gr.Column(scale=1):
+            output_gallery = gr.Gallery(
+                label="Imágenes similares", 
+                columns=3,
+                height="auto"
+            )
+    
+    search_button.click(
+        fn=search_text,
+        inputs=[input_text, top_k],
+        outputs=output_gallery
+    )
+
+if __name__ == "__main__":
+    from multiprocessing import freeze_support
+    freeze_support()
+    demo.launch()

model.py

@@ -0,0 +1,75 @@
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+
+class TextEncoderHead(nn.Module):
+    def __init__(self, model):
+        super(TextEncoderHead, self).__init__()
+        self.model = model
+        for param in self.model.parameters():
+            param.requires_grad = False
+        # uncomment this for chemberta
+        # self.seq1 = nn.Sequential(
+        #     nn.Flatten(),
+        #     nn.Linear(767*256, 2000),
+        #     nn.Dropout(0.3),
+        #     nn.ReLU(),
+        #     nn.Linear(2000, 512),
+        #     nn.LayerNorm(512)
+        # )
+        self.seq1 = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(768*256, 2000),
+            nn.Dropout(0.3),
+            nn.ReLU(),
+            nn.Linear(2000, 512),
+            nn.LayerNorm(512)
+        )
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        # uncomment this for chemberta
+        # outputs = outputs.logits
+        outputs = outputs.last_hidden_state
+        outputs = self.seq1(outputs)
+        return outputs.contiguous()
+    
+class ImageEncoderHead(nn.Module):
+    def __init__(self, model):
+        super(ImageEncoderHead, self).__init__()
+        self.model = model
+        for param in self.model.parameters():
+            param.requires_grad = False
+        # for resnet model
+        # self.seq1 = nn.Sequential(
+        #     nn.Flatten(),
+        #     nn.Linear(512*7*7, 1000),
+        #     nn.Linear(1000, 512),
+        #     nn.LayerNorm(512)
+        # )
+        # for vit model
+        self.seq1 = nn.Sequential(
+            nn.Linear(768, 1000),
+            nn.Dropout(0.3),
+            nn.ReLU(),
+            nn.Linear(1000, 512),
+            nn.LayerNorm(512)
+        )
+
+    
+    def forward(self, pixel_values):
+        outputs = self.model(pixel_values)
+        outputs = outputs.last_hidden_state.mean(dim=1)
+        outputs = self.seq1(outputs)
+        return outputs.contiguous()
+    
+class CLIPChemistryModel(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, text_encoder, image_encoder):
+        super(CLIPChemistryModel, self).__init__()
+        self.text_encoder = text_encoder
+        self.image_encoder = image_encoder
+
+    def forward(self, image, input_ids, attention_mask):
+        # calculate the embeddings
+        ie = self.image_encoder(image)
+        te = self.text_encoder(input_ids, attention_mask)
+        return ie, te

requirements.txt

@@ -0,0 +1,3 @@
+pinecone=="5.4.2"
+transformers=="4.47.0"
+huggingface-hub=="0.26.5"

search.py

@@ -0,0 +1,105 @@
+from transformers import ViTModel, AutoModelForMaskedLM, AutoTokenizer, ViTImageProcessor, DistilBertModel
+from pinecone import Pinecone
+import torch
+
+
+pc = Pinecone()
+index = pc.Index("clipmodel")
+
+
+from io import BytesIO
+import base64
+from PIL import Image
+
+import sys
+
+sys.path.append('../src')
+
+from model import CLIPChemistryModel, TextEncoderHead, ImageEncoderHead
+
+
+ENCODER_BASE = DistilBertModel.from_pretrained("distilbert-base-uncased")
+IMAGE_BASE = ViTModel.from_pretrained("google/vit-base-patch16-224")
+text_encoder = TextEncoderHead(model=ENCODER_BASE)
+image_encoder = ImageEncoderHead(model=IMAGE_BASE)
+
+clip_model = CLIPChemistryModel(text_encoder=text_encoder, image_encoder=image_encoder)
+
+clip_model.load_state_dict(torch.load('/Users/sebastianalejandrosarastizambonino/Documents/projects/CLIP_Pytorch/src/best_model_fashion.pth', map_location=torch.device('cpu')))
+
+te_final = clip_model.text_encoder
+ie_final = clip_model.image_encoder
+
+def process_text_for_encoder(text, model):
+    # tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
+    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+    encoded_input = tokenizer(text, return_tensors='pt', padding='max_length', truncation=True, max_length=256)
+    input_ids = encoded_input['input_ids']
+    attention_mask = encoded_input['attention_mask']
+    output = model(input_ids=input_ids, attention_mask=attention_mask)
+    return output.detach().numpy().tolist()[0]
+
+def process_image_for_encoder(image, model):
+    # image = Image.open(BytesIO(image))
+    print(type(image))
+    image_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
+    image_tensor = image_processor(image, 
+            return_tensors="pt", 
+            do_resize=True
+            )['pixel_values']
+    output =  model(pixel_values=image_tensor)
+    return output.detach().numpy().tolist()[0]
+
+def search_similarity(input, mode, top_k=5):
+    if mode == 'text':
+        output = process_text_for_encoder(input, model=te_final)
+    else:
+        output = input
+    
+    if mode == 'text':
+        mode_search = 'image'
+        response = index.query(
+            namespace="space-" + mode_search + "-fashion",
+            vector=output,
+            top_k=top_k,
+            include_values=True,
+            include_metadata=True
+        )
+        similar_images = [value['metadata']['image'] for value in response['matches']]
+        return similar_images
+    elif mode == 'image':
+        mode_search = 'text'
+        response = index.query(
+            namespace="space-" + mode_search + "-fashion",
+            vector=output,
+            top_k=top_k,
+            include_values=True,
+            include_metadata=True
+        )
+        similar_text = [value['metadata']['text'] for value in response['matches']]
+        return similar_text
+    else:
+        raise ValueError("mode must be either 'text' or 'image'")
+    
+def process_image_for_encoder_gradio(image, is_bytes=True):
+    """Procesa tanto imágenes en bytes como objetos PIL Image"""
+    try:
+        if is_bytes:
+            # Si la imagen viene en bytes
+            image = Image.open(BytesIO(image))
+        else:
+            # Si la imagen ya es un objeto PIL Image o viene de gradio
+            if not isinstance(image, Image.Image):
+                # Si viene de gradio, podría ser un numpy array
+                image = Image.fromarray(image)
+        
+        image_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
+        image_tensor = image_processor(image, 
+                return_tensors="pt", 
+                do_resize=True
+                )['pixel_values']
+        output = ie_final(pixel_values=image_tensor)
+        return output.detach().numpy().tolist()[0]
+    except Exception as e:
+        print(f"Error en process_image_for_encoder: {e}")
+        raise

utils.py

@@ -0,0 +1,34 @@
+from model import CLIPChemistryModel, TextEncoderHead, ImageEncoderHead
+from transformers import ViTModel, AutoModelForMaskedLM, AutoTokenizer, ViTImageProcessor
+
+from io import BytesIO
+import base64
+
+from PIL import Image
+
+def bytes_to_str(bytes_data):
+    return base64.b64encode(bytes_data).decode('utf-8')
+
+def str_to_bytes(str_data):
+    return base64.b64decode(str_data)
+
+def push_embeddings_to_pine_cone(index, embeddings, df, mode, length):
+    records = []
+    for i in range(length):
+        if mode == 'text':
+            records.append({
+                "id": str(mode) + str(i), 
+                "values": embeddings[i],
+                "metadata": {str(mode): df[mode].iloc[i]}})
+        elif mode == 'image':
+            records.append({
+                "id": str(mode) + str(i), 
+                "values": embeddings[i],
+                "metadata": {str(mode): bytes_to_str(df[mode].iloc[i]['bytes'])}})
+        else:
+            raise ValueError("mode must be either 'text' or 'image'")
+    
+    index.upsert(
+        vectors=records,
+        namespace="space-" + mode
+    )