add imgs

app.py

@@ -10,6 +10,7 @@ from imagebind import data
 from imagebind.models import imagebind_model
 from imagebind.models.imagebind_model import ModalityType
 import torch.nn as nn
+import pickle
 
 
 device = "cpu" #"cuda:0" if torch.cuda.is_available() else "cpu"
@@ -17,9 +18,29 @@ model = imagebind_model.imagebind_huge(pretrained=True)
 model.eval()
 model.to(device)
 
+image_features = pickle.load(open("./assets/image_features.pkl","rb"))
+image_paths = pickle.load(open("./assets/image_paths.pkl","rb"))
+
 
 def generate_image(text):
-    return Image.open("./assets/ICA-Logo.png").convert("RGB")
+    inputs = {
+        ModalityType.TEXT: data.load_and_transform_text([text], device)
+    }
+
+    with torch.no_grad():
+        embeddings = model(inputs)
+    
+    text_features = embeddings[ModalityType.TEXT]
+    text_features /= text_features.norm(dim=-1, keepdim=True)
+
+    similarity = text_features.cpu().numpy() @ image_features.cpu().numpy().T
+
+    #pega index maior
+    index_img = np.argmax(similarity)
+    img_name = os.path.basename(image_paths[index_img])
+    im = Image.open(f"./assets/images/{img_name}").convert("RGB")
+
+    return im
 
 # Interface do Gradio
 iface = gr.Interface(

assets/image_paths.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:512b07acbc5f8a2b02a78e8dbb417f86ee68fa87906101aa9b15c527b437b818
+size 435

requirements.txt

@@ -13,4 +13,5 @@ numpy>=1.19
 matplotlib
 types-regex
 mayavi
-scikit-image
+scikit-image
+pickle

utils.py

@@ -1,5 +1,7 @@
 import os
+import pickle
 
+# teste = pickle.load(open("df_to_search.pkl","rb"))
 def calculate_similarity():
     image_features = embeddings[ModalityType.VISION]
     text_features = embeddings[ModalityType.TEXT]