update the code files, adding flicker 8k dataset

app.py

@@ -2,13 +2,18 @@ import gradio as gr
 import utils
 
 
+
 # Araclip demo 
 with gr.Blocks() as demo_araclip:
 
-    gr.Markdown("## Input parameters")
+    gr.Markdown("## Choose the dataset")
 
+    dadtaset_select = gr.Radio(["XTD dataset", "Flicker 8k dataset"], value="XTD dataset", label="Dataset", info="Which dataset you would like to search in?")
+
+    gr.Markdown("## Input parameters")
+    
     txt = gr.Textbox(label="Text Query (Caption)")
-    num = gr.Slider(label="Number of retrieved image", value=1, minimum=1, maximum=1000)
+    num = gr.Slider(label="Number of retrieved image", value=1, minimum=1)
     
 
     with gr.Row():
@@ -32,20 +37,18 @@ with gr.Blocks() as demo_araclip:
             json_output = gr.JSON()
 
         with gr.Column(scale=1):
-            # gr.Markdown("### Data Retrieved based on Text similarity")
-            # gr.Markdown("<div style='text-align: center;'> Data Retrieved based on Text similarity </div>")
             gr.Markdown("<div style='text-align: center; font-size: 24px; font-weight: bold;'>Data Retrieved based on Text similarity</div>")
             json_text = gr.JSON()
 
 
-    btn.click(utils.predict, inputs=[txt, num], outputs=[gallery,lables, json_output, json_text])
+    btn.click(utils.predict, inputs=[txt, num, dadtaset_select], outputs=[gallery,lables, json_output, json_text])
 
 
     gr.Examples(
         examples=[["تخطي لاعب فريق بيتسبرج بايرتس منطقة اللوحة الرئيسية في مباراة بدوري البيسبول", 5], 
                   ["وقوف قطة بمخالبها على فأرة حاسوب على المكتب", 10],
                   ["صحن به شوربة صينية بالخضار، وإلى جانبه بطاطس مقلية وزجاجة ماء", 7]],
-        inputs=[txt, num],
+        inputs=[txt, num, dadtaset_select],
         outputs=[gallery,lables, json_output, json_text],
         fn=utils.predict,
         cache_examples=False,
@@ -54,10 +57,15 @@ with gr.Blocks() as demo_araclip:
 # mclip demo 
 with gr.Blocks() as demo_mclip:
 
+    gr.Markdown("## Choose the dataset")
+
+    dadtaset_select = gr.Radio(["XTD dataset", "Flicker 8k dataset"], value="XTD dataset", label="Dataset", info="Which dataset you would like to search in?")
+
+
     gr.Markdown("## Input parameters")
     
     txt = gr.Textbox(label="Text Query (Caption)")
-    num = gr.Slider(label="Number of retrieved image", value=1, minimum=1, maximum=1000)
+    num = gr.Slider(label="Number of retrieved image", value=1, minimum=1)
 
     with gr.Row():
         btn = gr.Button("Retrieve images", scale=1)
@@ -81,16 +89,13 @@ with gr.Blocks() as demo_mclip:
             gr.Markdown("## Text Retrieved")
             json_text = gr.JSON()
 
-    btn.click(utils.predict_mclip, inputs=[txt, num], outputs=[gallery,lables, json_output, json_text])
-
-    
-
+    btn.click(utils.predict_mclip, inputs=[txt, num, dadtaset_select], outputs=[gallery,lables, json_output, json_text])
 
     gr.Examples(
         examples=[["تخطي لاعب فريق بيتسبرج بايرتس منطقة اللوحة الرئيسية في مباراة بدوري البيسبول", 5], 
                   ["وقوف قطة بمخالبها على فأرة حاسوب على المكتب", 10],
                   ["صحن به شوربة صينية بالخضار، وإلى جانبه بطاطس مقلية وزجاجة ماء", 7]],
-        inputs=[txt, num],
+        inputs=[txt, num, dadtaset_select],
         outputs=[gallery,lables, json_output, json_text],
         fn=utils.predict_mclip,
         cache_examples=False,
@@ -101,8 +106,8 @@ with gr.Blocks() as demo_mclip:
 with gr.Blocks() as demo:
 
     gr.Markdown("<font color=red size=10><center>AraClip: Arabic Image Retrieval Application</center></font>")
-    gr.TabbedInterface([demo_araclip, demo_mclip], ["Our Model", "Mclip model"])
 
+    gr.TabbedInterface([demo_araclip, demo_mclip], ["Our Model", "Mclip model"])
 
 
 if __name__ == "__main__":

utils.py

@@ -43,10 +43,9 @@ def features_pickle(file_path=None):
     return features_pickle
 
 
-def dataset_loading():
-
-    with open("photos/en_ar_XTD10_edited_v2.jsonl") as filino:
+def dataset_loading(file_name):
 
+    with open(file_name) as filino:
 
         data = [json.loads(file_i) for file_i in filino]
 
@@ -89,7 +88,7 @@ def compare_embeddings_text(full_text_embds, txt_embs):
 
 
 
-def find_image(language_model,clip_model, text_query, dataset, image_features, text_features_new,sorted_data, num=1):
+def find_image(language_model,clip_model, text_query, dataset, image_features, text_features_new,sorted_data, images_path,num=1):
 
     embedding, _  = text_encoder(language_model, text_query)
 
@@ -111,7 +110,7 @@ def find_image(language_model,clip_model, text_query, dataset, image_features, t
 
         for i in range(1, num+1):
             idx = np.argsort(probs, axis=0)[-i, 0]
-            path = 'photos/XTD10_dataset/' + dataset.get_image_name(idx)
+            path = images_path + dataset.get_image_name(idx)
                     
             path_l = (path,f"{sorted_data[idx]['caption_ar']}")
 
@@ -142,27 +141,32 @@ class AraClip():
         self.text_model = load_model('bert-base-arabertv2-ViT-B-16-SigLIP-512-epoch-155-trained-2M', in_features= 768, out_features=768)
         self.language_model = lambda queries: np.asarray(self.text_model(queries).detach().to('cpu')) 
         self.clip_model, self.compose = create_model_from_pretrained('hf-hub:timm/ViT-B-16-SigLIP-512')
-        self.sorted_data, self.image_name_list = dataset_loading()
 
-    def load_images(self):
-        # Return the features of the text and images
-        image_features_new = features_pickle('cashed_pickles/image_features_XTD_1000_images_arabert_siglib_best_model.pickle')
-        return image_features_new
-    
-    def load_text(self):
-        text_features_new = features_pickle('cashed_pickles/text_features_XTD_1000_images_arabert_siglib_best_model.pickle')
-        return text_features_new
-    
-    def load_dataset(self):
-        dataset = CustomDataSet("photos/XTD10_dataset", self.compose, self.image_name_list)
+        self.sorted_data_xtd, self.image_name_list_xtd = dataset_loading("photos/en_ar_XTD10_edited_v2.jsonl")
+        self.sorted_data_flicker8k, self.image_name_list_flicker8k = dataset_loading("photos/flicker_8k.jsonl")
+
+    def load_pickle_file(self, file_name):
+
+        return features_pickle(file_name) 
+
+    def load_xtd_dataset(self):
+        dataset = CustomDataSet("photos/XTD10_dataset", self.compose, self.image_name_list_xtd)
+        
         return dataset
 
+    def load_flicker8k_dataset(self):
+        dataset = CustomDataSet("photos/Flicker8k_Dataset", self.compose, self.image_name_list_flicker8k)
+        return dataset
 
 araclip = AraClip()
 
-def predict(text, num):
+def predict(text, num, dadtaset_select):
+
+    if dadtaset_select == "XTD dataset":
+        image_paths, labels, json_data, json_text = find_image(araclip.language_model,araclip.clip_model, text, araclip.load_xtd_dataset(), araclip.load_pickle_file("cashed_pickles/XTD_pickles/image_features_XTD_1000_images_arabert_siglib_best_model.pickle") , araclip.load_pickle_file("cashed_pickles/XTD_pickles/image_features_XTD_1000_images_arabert_siglib_best_model.pickle"), araclip.sorted_data_xtd, 'photos/XTD10_dataset/', num=int(num))
 
-    image_paths, labels, json_data, json_text = find_image(araclip.language_model,araclip.clip_model, text, araclip.load_dataset(), araclip.load_images() , araclip.load_text(), araclip.sorted_data, num=int(num))
+    else:
+        image_paths, labels, json_data, json_text = find_image(araclip.language_model,araclip.clip_model, text, araclip.load_flicker8k_dataset(), araclip.load_pickle_file("cashed_pickles/flicker_8k/image_features_flicker_8k_images_arabert_siglib_best_model.pickle") , araclip.load_pickle_file("cashed_pickles/flicker_8k/text_features_flicker_8k_images_arabert_siglib_best_model.pickle"), araclip.sorted_data_flicker8k, "photos/Flicker8k_Dataset/", num=int(num))
 
     return image_paths, labels, json_data, json_text
 
@@ -175,26 +179,34 @@ class Mclip():
         self.text_model_mclip = pt_multilingual_clip.MultilingualCLIP.from_pretrained('M-CLIP/XLM-Roberta-Large-Vit-B-16Plus')
         self.language_model_mclip = lambda queries:  np.asarray(self.text_model_mclip.forward(queries, self.tokenizer_mclip).detach().to('cpu'))  
         self.clip_model_mclip, _, self.compose_mclip = create_model_and_transforms('ViT-B-16-plus-240', pretrained="laion400m_e32")
-        self.sorted_data, self.image_name_list = dataset_loading()
+        self.sorted_data_xtd, self.image_name_list_xtd = dataset_loading("photos/en_ar_XTD10_edited_v2.jsonl")
+        self.sorted_data_flicker8k, self.image_name_list_flicker8k = dataset_loading("photos/flicker_8k.jsonl")
 
-    def load_images(self):
-        # Return the features of the text and images
-        image_features_mclip = features_pickle('cashed_pickles/image_features_XTD_1000_images_XLM_Roberta_Large_Vit_B_16Plus_ar.pickle')
-        return image_features_mclip
-    
-    def load_text(self):
-        text_features_new_mclip = features_pickle('cashed_pickles/text_features_XTD_1000_images_XLM_Roberta_Large_Vit_B_16Plus_ar.pickle')
-        return text_features_new_mclip
-    
-    def load_dataset(self):
-        dataset_mclip = CustomDataSet("photos/XTD10_dataset", self.compose_mclip, self.image_name_list)
-        return dataset_mclip
+    def load_pickle_file(self, file_name):
 
+        return features_pickle(file_name) 
+
+
+    def load_xtd_dataset(self):
+        dataset = CustomDataSet("photos/XTD10_dataset", self.compose_mclip, self.image_name_list_xtd)
+        
+        return dataset
+
+    def load_flicker8k_dataset(self):
+        dataset = CustomDataSet("photos/Flicker8k_Dataset", self.compose_mclip, self.image_name_list_flicker8k)
+        return dataset
+    
 
 mclip = Mclip()
 
-def predict_mclip(text, num):
+def predict_mclip(text, num, dadtaset_select):
+
 
-    image_paths, labels, json_data, json_text = find_image(mclip.language_model_mclip,mclip.clip_model_mclip, text, mclip.load_dataset() , mclip.load_text() , mclip.load_text() , mclip.sorted_data , num=int(num))
+    if dadtaset_select == "XTD dataset":
+        image_paths, labels, json_data, json_text = find_image(mclip.language_model_mclip,mclip.clip_model_mclip, text, mclip.load_xtd_dataset() , mclip.load_pickle_file("cashed_pickles/XTD_pickles/image_features_XTD_1000_images_XLM_Roberta_Large_Vit_B_16Plus_ar.pickle") , mclip.load_pickle_file("cashed_pickles/XTD_pickles/text_features_XTD_1000_images_XLM_Roberta_Large_Vit_B_16Plus_ar.pickle") , mclip.sorted_data_xtd , 'photos/XTD10_dataset/', num=int(num))
 
+    else:
+        image_paths, labels, json_data, json_text = find_image(mclip.language_model_mclip,mclip.clip_model_mclip, text, mclip.load_flicker8k_dataset() , mclip.load_pickle_file("cashed_pickles/flicker_8k/image_features_flicker_8k_images_XLM_Roberta_Large_Vit_B_16Plus_ar.pickle") , mclip.load_pickle_file("cashed_pickles/flicker_8k/text_features_flicker_8k_images_XLM_Roberta_Large_Vit_B_16Plus_ar.pickle") , mclip.sorted_data_flicker8k , 'photos/Flicker8k_Dataset/', num=int(num))
+
+    
     return image_paths, labels, json_data, json_text