S3BIR app demo

app.py

@@ -0,0 +1,271 @@
+import os
+
+import streamlit as st
+from io import BytesIO
+import base64
+from multiprocessing.dummy import Pool
+from PIL import Image, ImageDraw
+
+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 
+
+token = os.getenv("HUGGINGFACE_TOKEN")
+
+# Autentica usando el token
+login(token=token)
+
+# Variables
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+HEIGHT = 200
+N_RESULTS = 15
+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():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    path_images = 'data'
+
+    # Descargar el modelo desde Hugging Face
+    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_checkpoint = torch.load(path_model, map_location=device)
+    model.load_state_dict(model_checkpoint['state_dict'])
+    model.eval()
+    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")
+    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 + embeddings[0][i][1].split("/images")[-1])
+
+    for i in range(len(embeddings[1])):
+        embeddings[1][i] = (embeddings[1][i][0], path_images + embeddings[1][i][1].split("/images")[-1])
+
+    return model, path_images, embeddings
+
+
+def compute_text_embeddings(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)
+    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)
+
+    # 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()}
+    label_of_images = torch.tensor(
+        [path_to_label[item[1]] for item in embeddings[corpus_id]]).to(device)
+
+    return [
+        (
+            # path_images +  "page" + str(i) + ".jpg", # DocExplore
+            label_to_path[i],  # DocExplore
+        )
+        for i in label_of_images[max_indices].cpu().numpy().tolist()
+    ], dot_product[max_indices]  # bbox_of_images[max_indices], dot_product[max_indices]
+
+
+def make_square(img, fill_color=(255, 255, 255)):
+    x, y = img.size
+    size = max(x, y)
+    new_img = Image.new("RGB", (x, y), fill_color)
+    new_img.paste(img)
+    return new_img, x, y
+
+
+@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
+
+
+def convert_pil_to_base64(image):
+    img_buffer = BytesIO()
+    image.save(img_buffer, format="JPEG")
+    byte_data = img_buffer.getvalue()
+    base64_str = base64.b64encode(byte_data)
+    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)):
+        title, encoded = url_list[i][0], encoded_images[i]
+        html = (
+            html
+            + f"<img title='{escape(title)}' style='height: {HEIGHT}px; margin: 5px' src='data:image/jpeg;base64,{encoded.decode()}'>"
+        )
+    html += "</div>"
+    return html
+
+
+description = """
+#  Sketch-based Image Retrieval (SBIR)
+"""
+
+div_style = {
+    "display": "flex",
+    "justify-content": "center",
+    "flex-wrap": "wrap",
+}
+
+
+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():
+    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;
+        }
+        </style>
+        """,
+        unsafe_allow_html=True,
+    )
+    st.sidebar.markdown(description)
+
+    st.title("SBIR App")
+    _, col, _ = st.columns((1, 1, 1))
+    with col:
+        canvas_result = st_canvas(
+            background_color="#eee",
+            stroke_width=stroke_width,
+            update_streamlit=True,
+            height=300,
+            width=300,
+            key="color_annotation_app",
+        )
+
+    _, c, _ = st.columns((1, 3, 1))
+    query = ["koala"]  # c.text_input("", value="koala")
+    corpus = c.radio("", ["Ecommerce"])
+
+    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)
+        draw_tensor = transforms.Normalize(
+            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)
+        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)
+
+
+if __name__ == "__main__":
+    main()