Se actualiza el dataset

app.py

@@ -4,29 +4,29 @@ import streamlit as st
 from io import BytesIO
 import base64
 from multiprocessing.dummy import Pool
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageOps
 
 import torch
 from torchvision import transforms
 
 # sketches
 from streamlit_drawable_canvas import st_canvas
-from PIL import Image, ImageOps
-from torchvision import transforms
 from src.model_LN_prompt import Model
 
 
 import pickle as pkl
 from html import escape
-from huggingface_hub import hf_hub_download,login 
+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)
+login(token=token, add_to_git_credential=True)
 
 # Variables
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+print(f"Device: {device}")
 HEIGHT = 200
 N_RESULTS = 15
 color = st.get_option("theme.primaryColor")
@@ -35,25 +35,30 @@ if color is None:
 else:
     color = tuple(int(color.lstrip("#")[i: i + 2], 16) for i in (0, 2, 4))
 
-
 @st.cache_resource
 def load():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    path_images = 'data'
+    print("Cargando todo...")
+    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
-    path_model = hf_hub_download(repo_id="CHSTR/Ecommerce", filename="dinov2_ecommerce.ckpt")
+    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().to(device)
+    model = Model()
     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
-    embeddings_file = hf_hub_download(repo_id="CHSTR/Ecommerce", filename="ecommerce_demo.pkl")
+    embeddings_file = hf_hub_download(
+        repo_id="CHSTR/Ecommerce", filename="ecommerce_demo.pkl")
     print(f"Archivo de embeddings descargado en: {embeddings_file}")
 
     embeddings = {
@@ -63,26 +68,27 @@ def load():
 
     # 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 + embeddings[0][i][1].split("/images")[-1])
+        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 + embeddings[1][i][1].split("/images")[-1])
+        embeddings[1][i] = (embeddings[1][i][0], path_images +
+                            "/".join(embeddings[1][i][1].split("/")[-3:]))
 
     return model, path_images, embeddings
 
-
-def compute_text_embeddings(sketch):
+def compute_sketch(sketch):
     with torch.no_grad():
         sketch_feat = model(sketch.to(device), dtype='sketch')
     return sketch_feat
 
-
 def image_search(query, corpus, n_results=N_RESULTS):
-    query_embedding = compute_text_embeddings(query)
+    query_embedding = compute_sketch(query)
     corpus_id = 0 if corpus == "Unsplash" else 1
     image_features = torch.tensor(
         [item[0] for item in embeddings[corpus_id]]).to(device)
-    
+
     dot_product = (image_features @ query_embedding.T)[:, 0]
     _, max_indices = torch.topk(
         dot_product, n_results, dim=0, largest=True, sorted=True)
@@ -96,11 +102,10 @@ def image_search(query, corpus, n_results=N_RESULTS):
 
     return [
         (
-            # path_images +  "page" + str(i) + ".jpg", # DocExplore
-            label_to_path[i],  # DocExplore
+            label_to_path[i],
         )
         for i in label_of_images[max_indices].cpu().numpy().tolist()
-    ], dot_product[max_indices]  # bbox_of_images[max_indices], dot_product[max_indices]
+    ], dot_product[max_indices] 
 
 
 def make_square(img, fill_color=(255, 255, 255)):
@@ -171,20 +176,14 @@ div_style = {
 }
 
 
-print("Cargando modelos...")
 model, path_images, embeddings = load()
-source = {0: "\Ecommerce", 1: "\nNone"}
 
-stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 5)
 
-dataset_transforms = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
+def main():
 
+    print("Cargando modelos...")
+    stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 5)
 
-def main():
     st.markdown(
         """
         <style>
@@ -234,9 +233,8 @@ def main():
             key="color_annotation_app",
         )
 
-    _, c, _ = st.columns((1, 3, 1))
-    query = ["koala"]  # c.text_input("", value="koala")
-    corpus = c.radio("", ["Ecommerce"])
+    st.columns((1, 3, 1))
+    corpus = ["Ecommerce"]
 
     if canvas_result.image_data is not None:
         draw = Image.fromarray(canvas_result.image_data.astype("uint8"))
@@ -249,11 +247,8 @@ def main():
             mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
         )(draw_tensor)
         draw_tensor = draw_tensor.unsqueeze(0)
-    else:
-        return
 
-    if len(query) > 0:
-        retrieved, dot_product = image_search(draw_tensor, corpus)
+        retrieved, _ = image_search(draw_tensor, corpus)
         imgs, xs, ys = get_images([x[0] for x in retrieved])
         encoded_images = []
         for image_idx in range(len(imgs)):
@@ -265,6 +260,8 @@ def main():
                 img0.resize((new_x, new_y))))
         st.markdown(get_html(retrieved, encoded_images),
                     unsafe_allow_html=True)
+    else:
+        return
 
 
 if __name__ == "__main__":