Initial Release

app.py

@@ -12,6 +12,8 @@ print(f'gr version : {gr.__version__}')
 import pickle
 import random
 
+import numpy as np
+
 # %%
 model_name = 'trclip-vitl14-e10'
 if not os.path.exists(model_name):
@@ -28,24 +30,37 @@ if not os.path.exists('TrCaption-trclip-vitl14-e10'):
 
 # %%
 
-def load_image_embeddings():
+def load_image_embeddings(load_batch=True):
     path = os.path.join('TrCaption-trclip-vitl14-e10', 'image_embeddings')
     bs = 100_000
-    embeddings = []
+    if load_batch:
+        for i in tqdm(range(0, 3_100_000, bs), desc='Loading TrCaption Image embeddings'):
+            with open(os.path.join(path, f'image_em_{i}.pkl'), 'rb') as f:
+                yield pickle.load(f)
+        return
+
+    else:
+        embeddings = []
+        for i in tqdm(range(0, 3_100_000, bs), desc='Loading TrCaption Image embeddings'):
+            with open(os.path.join(path, f'image_em_{i}.pkl'), 'rb') as f:
+                embeddings.append(pickle.load(f))
+        return torch.cat(embeddings, dim=0)
 
-    for i in tqdm(range(0, 3_100_000, bs), desc='Loading TrCaption Image embeddings'):
-        with open(os.path.join(path, f'image_em_{i}.pkl'), 'rb') as f:
-            embeddings.append(pickle.load(f))
-    return torch.cat(embeddings, dim=0)
 
-def load_text_embeddings():
+def load_text_embeddings(load_batch=True):
     path = os.path.join('TrCaption-trclip-vitl14-e10', 'text_embeddings')
     bs = 100_000
-    embeddings = []
-    for i in tqdm(range(0, 3_600_000, bs), desc='Loading TrCaption text embeddings'):
-        with open(os.path.join(path, f'text_em_{i}.pkl'), 'rb') as f:
-            embeddings.append(pickle.load(f))
-    return torch.cat(embeddings, dim=0)
+    if load_batch:
+        for i in tqdm(range(0, 3_600_000, bs), desc='Loading TrCaption text embeddings'):
+            with open(os.path.join(path, f'text_em_{i}.pkl'), 'rb') as f:
+                yield pickle.load(f)
+        return
+    else:
+        embeddings = []
+        for i in tqdm(range(0, 3_600_000, bs), desc='Loading TrCaption text embeddings'):
+            with open(os.path.join(path, f'text_em_{i}.pkl'), 'rb') as f:
+                embeddings.append(pickle.load(f))
+        return torch.cat(embeddings, dim=0)
 
 
 def load_metadata():
@@ -56,61 +71,64 @@ def load_metadata():
     trcap_urls = metadata['image_urls']
     return trcap_texts, trcap_urls
 
-def load_spesific_tensor(index , type , bs= 100_000):
+
+def load_spesific_tensor(index, type, bs=100_000):
     part = index // bs
     idx = index % bs
-    with open(os.path.join('TrCaption-trclip-vitl14-e10', f'{type}_embeddings', f'{type}_em_{part*bs}.pkl'), 'rb') as f:
+    with open(os.path.join('TrCaption-trclip-vitl14-e10', f'{type}_embeddings', f'{type}_em_{part * bs}.pkl'), 'rb') as f:
         embeddings = pickle.load(f)
     return embeddings[idx]
 
-# %%
 
-image_embeddings = None
-text_embeddings = None
-
-#%%
+# %%
 trcap_texts, trcap_urls = load_metadata()
 # %%
+print(f'INFO : Model loading')
 model_path = os.path.join(model_name, 'pytorch_model.bin')
 trclip = Trclip(model_path, clip_model='ViT-L/14', device='cpu')
-#%%
-import psutil
-
-print(f"First used memory {psutil.virtual_memory().used/float(1<<30):,.0f} GB" , )
 # %%
 
+
+
+# %%
 def run_im(im1, use_trcap_images, text1, use_trcap_texts):
+    print(f'INFO : Image retrieval starting')
     f_texts_embeddings = None
-    f_image_embeddings = None
-    global image_embeddings
-    global text_embeddings
     ims = None
-    print("im2", use_trcap_images)
     if use_trcap_images:
-        print('TRCaption images used')
-        # Images taken from TRCAPTION
+        print('INFO : TRCaption images used')
         im_paths = trcap_urls
-        if image_embeddings is None:
-            print(f"First used memory {psutil.virtual_memory().used / float(1 << 30):,.0f} GB", )
-            text_embeddings = None
-            image_embeddings = load_image_embeddings()
-            print(f"First used memory {psutil.virtual_memory().used / float(1 << 30):,.0f} GB", )
-        f_image_embeddings = image_embeddings
     else:
+        print('INFO : Own images used')
         # Images taken from user
         im_paths = [i.name for i in im1]
         ims = [Image.open(i) for i in im_paths]
     if use_trcap_texts:
+        print(f'INFO : TRCaption texts used')
         random_indexes = random.sample(range(len(trcap_texts)), 2)  # MAX 2 text are allowed in image retrieval UI limit
         f_texts_embeddings = []
         for i in random_indexes:
             f_texts_embeddings.append(load_spesific_tensor(i, 'text'))
         f_texts_embeddings = torch.stack(f_texts_embeddings)
         texts = [trcap_texts[i] for i in random_indexes]
+
     else:
-        texts = [i.trim() for i in text1.split('\n')[:2] if i.trim() != '']
+        print(f'INFO : Own texts used')
+        texts = [i.strip() for i in text1.split('\n')[:2] if i.strip() != '']
+
+    if use_trcap_images:  # This means that we will iterate over batches because Huggingface space has 16 gb limit :///
+        per_mode_probs = []
+        f_texts_embeddings = f_texts_embeddings if use_trcap_texts else trclip.get_text_features(texts)
+        for f_image_embeddings in tqdm(load_image_embeddings(load_batch=True), desc='Running image retrieval'):
+            batch_probs = trclip.get_results(
+                text_features=f_texts_embeddings, image_features=f_image_embeddings, mode='per_text', return_probs=True)
+            per_mode_probs.append(batch_probs)
+        per_mode_probs = torch.cat(per_mode_probs, dim=1)
+        per_mode_probs = per_mode_probs.softmax(dim=-1).cpu().detach().numpy()
+        per_mode_indices = [np.argsort(prob)[::-1] for prob in per_mode_probs]
 
-    per_mode_indices, per_mode_probs = trclip.get_results(texts=texts, images=ims, text_features=f_texts_embeddings, image_features=f_image_embeddings, mode='per_text')
+    else:
+        per_mode_indices, per_mode_probs = trclip.get_results(texts=texts, images=ims, text_features=f_texts_embeddings, mode='per_text')
 
     print(f'per_mode_indices = {per_mode_indices}\n,per_mode_probs = {per_mode_probs}  ')
     print(f'im_paths    = {im_paths}')
@@ -122,39 +140,45 @@ def run_im(im1, use_trcap_images, text1, use_trcap_texts):
 
 
 def run_text(im1, use_trcap_images, text1, use_trcap_texts):
-    f_texts_embeddings = None
+    print(f'INFO : Image retrieval starting')
     f_image_embeddings = None
-    global image_embeddings
-    global text_embeddings
     ims = None
     if use_trcap_images:
+        print('INFO : TRCaption images used')
         random_indexes = random.sample(range(len(trcap_urls)), 2)  # MAX 2 text are allowed in image retrieval UI limit
         f_image_embeddings = []
         for i in random_indexes:
             f_image_embeddings.append(load_spesific_tensor(i, 'image'))
         f_image_embeddings = torch.stack(f_image_embeddings)
-        print('TRCaption images used')
+        print(f'f_image_embeddings = {f_image_embeddings}')
         # Images taken from TRCAPTION
         im_paths = [trcap_urls[i] for i in random_indexes]
+        print(f'im_paths = {im_paths}')
+
     else:
+        print('INFO : Own images used')
         # Images taken from user
         im_paths = [i.name for i in im1[:2]]
         ims = [Image.open(i) for i in im_paths]
 
     if use_trcap_texts:
-        if text_embeddings is None:
-            print(f"Used memory {psutil.virtual_memory().used / float(1 << 30):,.0f} GB", )
-            image_embeddings = None
-            print(f"Image embd deleted used memory {psutil.virtual_memory().used / float(1 << 30):,.0f} GB", )
-            text_embeddings = load_text_embeddings()
-            print(f"Text embed  used memory {psutil.virtual_memory().used / float(1 << 30):,.0f} GB", )
-
-        f_texts_embeddings = text_embeddings
         texts = trcap_texts
     else:
-        texts = [i.trim() for i in text1.split('\n') if i.trim() != '']
+        texts = [i.strip() for i in text1.split('\n')[:2] if i.strip() != '']
+
+    if use_trcap_texts:
+        f_image_embeddings = f_image_embeddings if use_trcap_images else trclip.get_image_features(ims)
+        per_mode_probs = []
+        for f_texts_embeddings in tqdm(load_text_embeddings(load_batch=True), desc='Running text retrieval'):
+            batch_probs = trclip.get_results(
+                 text_features=f_texts_embeddings, image_features=f_image_embeddings, mode='per_image', return_probs=True)
+            per_mode_probs.append(batch_probs)
+        per_mode_probs = torch.cat(per_mode_probs, dim=1)
+        per_mode_probs = per_mode_probs.softmax(dim=-1).cpu().detach().numpy()
+        per_mode_indices = [np.argsort(prob)[::-1] for prob in per_mode_probs]
 
-    per_mode_indices, per_mode_probs = trclip.get_results(texts=texts, images=ims, image_features=f_image_embeddings, text_features=f_texts_embeddings, mode='per_image')
+    else:
+        per_mode_indices, per_mode_probs = trclip.get_results(texts=texts, images=ims, image_features=f_image_embeddings, mode='per_image')
     print(per_mode_indices)
     print(per_mode_probs)
     return text_retrieval_visualize(per_mode_indices, per_mode_probs, im_paths, texts,
@@ -219,7 +243,7 @@ with gr.Blocks() as demo:
                   <rect x="23" y="115" width="23" height="23" fill="#AEAEAE"></rect>
                   <rect x="23" y="69" width="23" height="23" fill="black"></rect>
                 </svg>
-                <h1 style="font-weight: 900; margin-bottom: 7px;">
+                <h1 style="font-weight: 1500; margin-bottom: 7px;">
                   Trclip Demo
                                   <a
                   href="https://github.com/yusufani/TrCLIP"
@@ -234,21 +258,35 @@ with gr.Blocks() as demo:
                 Also you can use pre calculated TrCaption embeddings. 
                 Number of texts  = 3533312
                 Number of images =  3070976
+                
+                Some images are not available in the internet because I downloaded and calculated TrCaption embeddings long time ago. Don't be suprise if you encounter with Image not found :D
+                
 
-                >
+                
               </p>
+              <p style="margin-bottom: 10px; font-size: 75%" ><em>Huggingface Space containers has 16 gb ram. TrCaption embeddings are totaly 20 gb. </em><em>I did a lot of writing and reading to files to make this space workable. That's why<span style="background-color: #ff6600; color: #ffffff;"> <strong>it's running much slower if you're using TrCaption Embeddig</strong>s</span>.</em></p>
+                <div class="sc-jSFjdj sc-iCoGMd jcTaHb kMthTr">
+                <div class="sc-iqAclL xfxEN">
+                <div class="sc-bdnxRM fJdnBK sc-crzoAE DykGo">
+                <div class="sc-gtsrHT gfuSqG">&nbsp;</div>
+                </div>
+                </div>
+                </div>
+                <div class="sc-jSFjdj sc-gKAaRy jcTaHb hydYaP">
+                <div class="sc-pNWdM lfZLSv">&nbsp;</div>
+                </div>
             </div>
         """)
 
     with gr.Tabs():
         with gr.TabItem("Use Own Images"):
             im_input = gr.components.File(label="Image input", optional=True, file_count='multiple')
-    is_trcap_ims = gr.Checkbox(label="Use TRCaption Images\nNote: ( Random 2 sample selected in text retrieval mode )")
+    is_trcap_ims = gr.Checkbox(label="Use TRCaption Images\n[Note: Random 2 sample selected in text retrieval mode )]")
 
     with gr.Tabs():
         with gr.TabItem("Input a text (Seperated by new line Max 2 for Image retrieval)"):
             text_input = gr.components.Textbox(label="Text input", optional=True)
-    is_trcap_texts = gr.Checkbox(label="Use TrCaption Captions \nNote: ( Random 2 sample selected in image retrieval mode")
+    is_trcap_texts = gr.Checkbox(label="Use TrCaption Captions \n[Note: Random 2 sample selected in image retrieval mode]")
 
     im_ret_but = gr.Button("Image Retrieval")
     text_ret_but = gr.Button("Text Retrieval")