Se utiliza el dataset desde hugginface

app.py

@@ -1,90 +1,85 @@
 import os
-
 import streamlit as st
 from io import BytesIO
-import base64
 from multiprocessing.dummy import Pool
-from PIL import Image, ImageDraw, ImageOps
-
+import base64
+from PIL import Image, ImageOps
 import torch
 from torchvision import transforms
-
-# sketches
 from streamlit_drawable_canvas import st_canvas
 from src.model_LN_prompt import Model
-
-
-import pickle as pkl
 from html import escape
+import pickle as pkl
 from huggingface_hub import hf_hub_download, login
 from datasets import load_dataset
 
-token = os.getenv("HUGGINGFACE_TOKEN")
 
-# Autentica usando el token
-login(token=token, add_to_git_credential=True)
+if 'initialized' not in st.session_state:
+    st.session_state.initialized = False
 
-# Variables
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-print(f"Device: {device}")
 HEIGHT = 200
-N_RESULTS = 15
+N_RESULTS = 20
 color = st.get_option("theme.primaryColor")
 if color is None:
     color = (0, 0, 255)
 else:
     color = tuple(int(color.lstrip("#")[i: i + 2], 16) for i in (0, 2, 4))
 
+
 @st.cache_resource
-def load():
-    print("Cargando todo...")
+def initialize_huggingface():
+    token = os.getenv("HUGGINGFACE_TOKEN")
+    if token:
+        login(token=token)
+    else:
+        st.error("HUGGINGFACE_TOKEN not found in environment variables")
+
+
+@st.cache_resource
+def load_model_and_data():
+    print("Loading everything...")
     dataset = load_dataset("CHSTR/ecommerce")
     path_images = "/".join(dataset['validation']
                            ['image'][0].filename.split("/")[:-3]) + "/"
-    print(f"Directorio de imágenes: {path_images}")
 
-    # Descargar el modelo desde Hugging Face
+    # Download model
     path_model = hf_hub_download(
         repo_id="CHSTR/Ecommerce", filename="dinov2_ecommerce.ckpt")
-    print(f"Archivo del modelo descargado en: {path_model}")
 
-    # Cargar el modelo
-    model = Model()
+    # Load model
+    model = Model().to(device)
     model_checkpoint = torch.load(path_model, map_location=device)
     model.load_state_dict(model_checkpoint['state_dict'])
     model.eval()
-    # model.to(device)
-    print("Modelo cargado exitosamente")
 
-    # Descargar y cargar los embeddings desde Hugging Face
+    # Download and load embeddings
     embeddings_file = hf_hub_download(
         repo_id="CHSTR/Ecommerce", filename="ecommerce_demo.pkl")
-    print(f"Archivo de embeddings descargado en: {embeddings_file}")
 
     embeddings = {
         0: pkl.load(open(embeddings_file, "rb")),
         1: pkl.load(open(embeddings_file, "rb"))
     }
 
-    # Actualizar los paths de las imágenes en los embeddings
-    for i in range(len(embeddings[0])):
-        embeddings[0][i] = (embeddings[0][i][0], path_images +
-                            "/".join(embeddings[0][i][1].split("/")[-3:]))
-        # print(embeddings[0][i])
-
-    for i in range(len(embeddings[1])):
-        embeddings[1][i] = (embeddings[1][i][0], path_images +
-                            "/".join(embeddings[1][i][1].split("/")[-3:]))
+    # Update image paths
+    for corpus_id in [0, 1]:
+        embeddings[corpus_id] = [
+            (emb[0], path_images + "/".join(emb[1].split("/")[-3:]))
+            for emb in embeddings[corpus_id]
+        ]
 
     return model, path_images, embeddings
 
-def compute_sketch(sketch):
+
+def compute_sketch(_sketch, model):
     with torch.no_grad():
-        sketch_feat = model(sketch.to(device), dtype='sketch')
+        sketch_feat = model(_sketch.to(device), dtype='sketch')
     return sketch_feat
 
-def image_search(query, corpus, n_results=N_RESULTS):
-    query_embedding = compute_sketch(query)
+
+def image_search(_query, corpus, model, embeddings, n_results=N_RESULTS):
+    query_embedding = compute_sketch(_query, model)
     corpus_id = 0 if corpus == "Unsplash" else 1
     image_features = torch.tensor(
         [item[0] for item in embeddings[corpus_id]]).to(device)
@@ -93,7 +88,6 @@ def image_search(query, corpus, n_results=N_RESULTS):
     _, max_indices = torch.topk(
         dot_product, n_results, dim=0, largest=True, sorted=True)
 
-    # Diccionario para mapear los paths a labels
     path_to_label = {path: idx for idx,
                      (_, path) in enumerate(embeddings[corpus_id])}
     label_to_path = {idx: path for path, idx in path_to_label.items()}
@@ -101,14 +95,14 @@ def image_search(query, corpus, n_results=N_RESULTS):
         [path_to_label[item[1]] for item in embeddings[corpus_id]]).to(device)
 
     return [
-        (
-            label_to_path[i],
-        )
+        (label_to_path[i],)
         for i in label_of_images[max_indices].cpu().numpy().tolist()
-    ], dot_product[max_indices] 
+    ], dot_product[max_indices]
 
 
-def make_square(img, fill_color=(255, 255, 255)):
+@st.cache_data
+def make_square(img_path, fill_color=(255, 255, 255)):
+    img = Image.open(img_path)
     x, y = img.size
     size = max(x, y)
     new_img = Image.new("RGB", (x, y), fill_color)
@@ -118,18 +112,12 @@ def make_square(img, fill_color=(255, 255, 255)):
 
 @st.cache_data
 def get_images(paths):
-    def process_image(path):
-        return make_square(Image.open(path))
-
-    processed = Pool(N_RESULTS).map(process_image, paths)
-    imgs, xs, ys = [], [], []
-    for img, x, y in processed:
-        imgs.append(img)
-        xs.append(x)
-        ys.append(y)
-    return imgs, xs, ys
+    processed = [make_square(path) for path in paths]
+    imgs, xs, ys = zip(*processed)
+    return list(imgs), list(xs), list(ys)
 
 
+@st.cache_data
 def convert_pil_to_base64(image):
     img_buffer = BytesIO()
     image.save(img_buffer, format="JPEG")
@@ -138,21 +126,6 @@ def convert_pil_to_base64(image):
     return base64_str
 
 
-def draw_reshape_encode(img, boxes, x, y):
-    boxes = [boxes.tolist()]
-    image = img.copy()
-    draw = ImageDraw.Draw(image)
-    new_x, new_y = int(x * HEIGHT / y), HEIGHT
-    for box in boxes:
-        print("box:", box)
-        draw.rectangle(
-            # (x_min, y_min, x_max, y_max)
-            [(box[0], box[1]), (box[2], box[3])],
-            outline=color,  # Box color
-            width=7  # Box width
-        )
-
-
 def get_html(url_list, encoded_images):
     html = "<div style='margin-top: 20px; max-width: 1200px; display: flex; flex-wrap: wrap; justify-content: space-evenly'>"
     for i in range(len(url_list)):
@@ -165,63 +138,40 @@ def get_html(url_list, encoded_images):
     return html
 
 
-description = """
-#  Sketch-based Image Retrieval (SBIR)
-"""
-
-div_style = {
-    "display": "flex",
-    "justify-content": "center",
-    "flex-wrap": "wrap",
-}
-
+def main():
+    if not st.session_state.initialized:
+        initialize_huggingface()
+        st.session_state.model, st.session_state.path_images, st.session_state.embeddings = load_model_and_data()
+        st.session_state.initialized = True
 
-model, path_images, embeddings = load()
+    description = """
+    #  Self-Supervised Sketch-based Image Retrieval (S3BIR)
 
+    Our approaches, S3BIR-CLIP and S3BIR-DINOv2, can produce a bimodal sketch-photo feature space from unpaired data without explicit sketch-photo pairs. Our experiments perform outstandingly in three diverse public datasets where the models are trained without real sketches.
 
-def main():
+    """
 
-    print("Cargando modelos...")
+    st.sidebar.markdown(description)
     stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 5)
 
+    # styles
     st.markdown(
         """
         <style>
-        .block-container{
-            max-width: 1200px;
-        }
-        div.row-widget > div{
-            flex-direction: row;
-            display: flex;
-            justify-content: center;
-        }
-        div.row-widget.stRadio > div > label{
-            margin-left: 5px;
-            margin-right: 5px;
-        }
-        .row-widget {
-            margin-top: -25px;
-        }
-        section > div:first-child {
-            padding-top: 30px;
-        }
-        div.appview-container > section:first-child{
-            max-width: 320px;
-        }
-        #MainMenu {
-            visibility: hidden;
-        }
-        .stMarkdown {
-            display: grid;
-            place-items: center;
-        }
+        .block-container{ max-width: 1200px; }
+        div.row-widget > div{ flex-direction: row; display: flex; justify-content: center; color: white; }
+        div.row-widget.stRadio > div > label{ margin-left: 5px; margin-right: 5px; }
+        .row-widget { margin-top: -25px; }
+        section > div:first-child { padding-top: 30px; }
+        div.appview-container > section:first-child{ max-width: 320px; }
+        #MainMenu { visibility: hidden; }
+        .stMarkdown { display: grid; place-items: center; }
         </style>
         """,
         unsafe_allow_html=True,
     )
-    st.sidebar.markdown(description)
 
-    st.title("SBIR App")
+    st.title("S3BIR App")
     _, col, _ = st.columns((1, 1, 1))
     with col:
         canvas_result = st_canvas(
@@ -233,13 +183,12 @@ def main():
             key="color_annotation_app",
         )
 
-    st.columns((1, 3, 1))
     corpus = ["Ecommerce"]
+    st.columns((1, 3, 1))
 
     if canvas_result.image_data is not None:
         draw = Image.fromarray(canvas_result.image_data.astype("uint8"))
         draw = ImageOps.pad(draw.convert("RGB"), size=(224, 224))
-        draw.save("draw.jpg")
 
         draw_tensor = transforms.ToTensor()(draw)
         draw_tensor = transforms.Resize((224, 224))(draw_tensor)
@@ -248,20 +197,19 @@ def main():
         )(draw_tensor)
         draw_tensor = draw_tensor.unsqueeze(0)
 
-        retrieved, _ = image_search(draw_tensor, corpus)
+        retrieved, _ = image_search(
+            draw_tensor, corpus[0], st.session_state.model, st.session_state.embeddings)
         imgs, xs, ys = get_images([x[0] for x in retrieved])
+
         encoded_images = []
         for image_idx in range(len(imgs)):
             img0, x, y = imgs[image_idx], xs[image_idx], ys[image_idx]
-
             new_x, new_y = int(x * HEIGHT / y), HEIGHT
-
             encoded_images.append(convert_pil_to_base64(
                 img0.resize((new_x, new_y))))
+
         st.markdown(get_html(retrieved, encoded_images),
                     unsafe_allow_html=True)
-    else:
-        return
 
 
 if __name__ == "__main__":

src/model_LN_prompt.py

@@ -1,15 +1,9 @@
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torchmetrics.functional import retrieval_average_precision
 import pytorch_lightning as pl
 
 from src.dinov2.models.vision_transformer import vit_base
-
-from functools import partial
-
-# from src.clip import clip
 from src.options import opts
 
 def freeze_model(m):
@@ -31,23 +25,11 @@ class Model(pl.LightningModule):
         self.opts = opts
 
         self.dino = vit_base(patch_size=14, block_chunks=0, init_values=1.0) 
-        print("self.dino", self.dino)
 
         # Prompt Engineering
         self.sk_prompt = nn.Parameter(torch.randn(self.opts.n_prompts, self.opts.prompt_dim))
         self.img_prompt = nn.Parameter(torch.randn(self.opts.n_prompts, self.opts.prompt_dim))
 
-        self.distance_fn = lambda x, y: 1.0 - F.cosine_similarity(x, y)
-        self.loss_fn_triplet = nn.TripletMarginWithDistanceLoss(
-             distance_function=self.distance_fn, margin=0.2)
-        
-        self.emb_cos_loss = nn.CosineEmbeddingLoss(margin=0.2)
-
-        self.loss_kl = nn.KLDivLoss(reduction="batchmean", log_target=True)
-
-        self.best_metric = -1e3
-        # normalization layer for the representations z1 and z2
-        # self.bn = nn.BatchNorm1d(self.opts.prompt_dim, affine=False)
 
     def configure_optimizers(self):
         if self.opts.model_type == 'one_encoder':

src/options.py

@@ -1,18 +1,17 @@
 import argparse
 
-parser = argparse.ArgumentParser(description='Sketch-based OD')
+parser = argparse.ArgumentParser(description='S3BIR')
 
-parser.add_argument('--exp_name', type=str, default='LN_prompt')
+parser.add_argument('--exp_name', type=str, default='DINOv2_prompt')
 
 # ----------------------
 # Training Params
 # ----------------------
 
-parser.add_argument('--clip_lr', type=float, default=1e-4)
-parser.add_argument('--clip_LN_lr', type=float, default=1e-6)
+parser.add_argument('--dinov2_lr', type=float, default=1e-4)
+parser.add_argument('--dinov2_LN_lr', type=float, default=1e-6)
 parser.add_argument('--prompt_lr', type=float, default=1e-4)
 parser.add_argument('--linear_lr', type=float, default=1e-4)
-parser.add_argument('--model_type', type=str, default='one_encoder', choices=['one_encoder', 'two_encoder'])
 
 # ----------------------
 # ViT Prompt Parameters