add class selection

app.py

@@ -14,11 +14,16 @@ num_nn = 20
 search_domain = 'all'
 num_results_per_domain = 5
 src_data_dict = {}
+class_list = []
 if search_domain == 'all':
     for d in domains:
         with open(os.path.join(feat_dir, f'dst_{d}_{shots}.pkl'), 'rb') as fp:
             src_data = pickle.load(fp)
-
+            if class_list == []:
+                for p in src_data[0]:
+                    cl = p.split('/')[-2]
+                    if cl not in class_list
+                        class_list.append(cl)
             src_nn_fit = NearestNeighbors(n_neighbors=num_results_per_domain, algorithm='auto', n_jobs=-1).fit(src_data[1])
             src_data_dict[d] = (src_data,src_nn_fit)
 else:
@@ -32,12 +37,17 @@ dst_data_dict = {}
 min_len = 1e10
 for d in domains:
     with open(os.path.join(feat_dir, f'src_{d}_{shots}.pkl'), 'rb') as fp:
-        dst_data_dict[d] = pickle.load(fp)
-        min_len = min(min_len, len(dst_data_dict[d][0]))
+        dest_data =  pickle.load(fp)
+        dst_data_dict[d] = {cl: [] for cl in class_list}
+        for p in dst_data[0]:
+            cl = p.split('/')[-2]
+            dst_data_dict[d][cl].append(p)
+        for cl in class_list:
+            min_len = min(min_len, len(dst_data_dict[d][cl]))
 
-def query(query_index, query_domain):
+def query(query_index, query_domain, cl):
     dst_data = dst_data_dict[query_domain]
-    dst_img_path = os.path.join(data_root, dst_data[0][query_index])
+    dst_img_path = os.path.join(data_root, dst_data[cl][query_index])
     img_paths = [dst_img_path]
     q_cl = dst_img_path.split('/')[-2]
     captions = [f'Query: {q_cl}'.title()]
@@ -61,17 +71,21 @@ with demo:
     gr.Markdown('## Instructions:')
     gr.Markdown('Select a query domain from the dropdown menu and the select any random image from the domain using the slider below. The retrieved results from each of the four domains, along with the class label will be presented.')
     gr.Markdown('## Select Query Domain: ')
-    domain_drop = gr.Dropdown(domains)
+    gr.Markdown('# Query Image: \t\t\t\t')
+    # domain_drop = gr.Dropdown(domains)
+    # cl_drop = gr.Dropdown(class_list)
 #     domain_select_button = gr.Button("Select Domain")
-    slider = gr.Slider(0, min_len)
+    # slider = gr.Slider(0, min_len)
     # slider = gr.Slider(0, 10000)
     image_button = gr.Button("Run")
-
     with gr.Row():
-        gr.Markdown('# Query Image: \t\t\t\t                   ')
-        gr.Markdown('\t')
-        gr.Markdown('\t')
-        gr.Markdown('\t')
+        with gr.Column():
+            domain_drop = gr.Dropdown(domains)
+            cl_drop = gr.Dropdown(class_list)
+            slider = gr.Slider(0, min_len)
+        # gr.Markdown('\t')
+        # gr.Markdown('\t')
+        # gr.Markdown('\t')
         with gr.Column():
             src_cap = gr.Label()
             src_img = gr.Image()
@@ -87,6 +101,6 @@ with demo:
                     out_captions.append(gr.Label())
                     out_images.append(gr.Image())
 
-    image_button.click(query, inputs=[slider, domain_drop], outputs=[src_img]+out_images +[src_cap]+ out_captions)
+    image_button.click(query, inputs=[slider, domain_drop, cl_drop], outputs=[src_img]+out_images +[src_cap]+ out_captions)
 
 demo.launch(share=True)