update UI

app.py

@@ -116,8 +116,12 @@ def dont_use_too_much_green(image_rgb):
 
 
 def to_pil_images(images):
+    size = images[0].shape[1]
+    target = 256
+    multiplier = target // size
+    res = int(size * multiplier)
     return [
-        Image.fromarray((image * 255).cpu().numpy().astype(np.uint8)).resize((256, 256), Image.Resampling.NEAREST)
+        Image.fromarray((image * 255).cpu().numpy().astype(np.uint8)).resize((res, res), Image.Resampling.NEAREST)
         for image in images
     ]
     
@@ -164,9 +168,10 @@ def get_random_path(length=10):
     path = f'/tmp/{name}.mp4'
     return path
 
-default_images = ['./images/image_0.jpg', './images/image_1.jpg', './images/image_2.jpg', './images/image_3.jpg', './images/image_5.jpg']
+default_images = ['./images/image_0.jpg', './images/image_1.jpg', './images/image_2.jpg', './images/image_3.jpg', './images/guitar_ego.jpg', './images/image_5.jpg']
 default_outputs = ['./images/image-1.webp', './images/image-2.webp', './images/image-3.webp', './images/image-4.webp', './images/image-5.webp']
-default_outputs_independent = ['./images/image-6.webp', './images/image-7.webp', './images/image-8.webp', './images/image-9.webp', './images/image-10.webp']
+# default_outputs_independent = ['./images/image-6.webp', './images/image-7.webp', './images/image-8.webp', './images/image-9.webp', './images/image-10.webp']
+default_outputs_independent = []
 
 downscaled_images = ['./images/image_0_small.jpg', './images/image_1_small.jpg', './images/image_2_small.jpg', './images/image_3_small.jpg', './images/image_5_small.jpg']
 downscaled_outputs = default_outputs
@@ -455,52 +460,50 @@ def run_fn(
 def make_input_images_section():
     gr.Markdown('### Input Images')
     input_gallery = gr.Gallery(value=None, label="Select images", show_label=False, elem_id="images", columns=[3], rows=[1], object_fit="contain", height="auto", type="pil", show_share_button=False)
-    submit_button = gr.Button("🔴 RUN", elem_id="submit_button")
-    clear_images_button = gr.Button("🗑️Clear", elem_id='clear_button')
+    submit_button = gr.Button("🔴 RUN", elem_id="submit_button", variant='primary')
+    clear_images_button = gr.Button("🗑️Clear", elem_id='clear_button', variant='stop')
     return input_gallery, submit_button, clear_images_button
 
 def make_input_video_section():
-    gr.Markdown('### Input Video')
+    # gr.Markdown('### Input Video')
     input_gallery = gr.Video(value=None, label="Select video", elem_id="video-input", height="auto", show_share_button=False)
     gr.Markdown('_image backbone model is used to extract features from each frame, NCUT is computed on all frames_')
     # max_frames_number = gr.Number(100, label="Max frames", elem_id="max_frames")
     max_frames_number = gr.Slider(1, 200, step=1, label="Max frames", value=100, elem_id="max_frames")
-    submit_button = gr.Button("🔴 RUN", elem_id="submit_button")
-    clear_images_button = gr.Button("🗑️Clear", elem_id='clear_button')
+    submit_button = gr.Button("🔴 RUN", elem_id="submit_button", variant='primary')
+    clear_images_button = gr.Button("🗑️Clear", elem_id='clear_button', variant='stop')
     return input_gallery, submit_button, clear_images_button, max_frames_number
 
-def make_example_images_section():
-    gr.Markdown('### Load Images 👇')
-    load_images_button = gr.Button("Load Example", elem_id="load-images-button")
-    example_gallery = gr.Gallery(value=example_items, label="Example Set A", show_label=False, columns=[3], rows=[2], object_fit="scale-down", height="200px", show_share_button=False, elem_id="example-gallery")
-    hide_button = gr.Button("Hide Example", elem_id="hide-button")
-    hide_button.click(
-        fn=lambda: gr.update(visible=False),
-        outputs=example_gallery
-    )
-    hide_button.click(
-        fn=lambda: gr.update(visible=False),
-        outputs=hide_button
-    )
-    return load_images_button, example_gallery, hide_button
+def make_dataset_images_section(advanced=False):
 
-def make_example_video_section():
-    gr.Markdown('### Load Video 👇')
-    load_video_button = gr.Button("Load Example", elem_id="load-video-button")
-    return load_video_button
-
-def make_dataset_images_section(open=False):
-    
-    with gr.Accordion("➡️ Click to expand: Load from dataset", open=open):
+    gr.Markdown('### Load Datasets')
+    load_images_button = gr.Button("Load", elem_id="load-images-button", variant='secondary')
+    advanced_radio = gr.Radio(["Basic", "Advanced"], label="Datasets", value="Advanced" if advanced else "Basic", elem_id="advanced-radio")
+    with gr.Column() as basic_block:
+        example_gallery = gr.Gallery(value=example_items, label="Example Set A", show_label=False, columns=[3], rows=[2], object_fit="scale-down", height="200px", show_share_button=False, elem_id="example-gallery")
+    with gr.Column() as advanced_block:
         dataset_names = DATASET_NAMES
         dataset_classes = DATASET_CLASSES
         dataset_dropdown = gr.Dropdown(dataset_names, label="Dataset name", value="mrm8488/ImageNet1K-val", elem_id="dataset")
         num_images_slider = gr.Number(10, label="Number of images", elem_id="num_images")
+        # with gr.Row():
         filter_by_class_checkbox = gr.Checkbox(label="Filter by class", value=True, elem_id="filter_by_class_checkbox")
         filter_by_class_text = gr.Textbox(label="Class to select", value="0,33,99", elem_id="filter_by_class_text", info=f"e.g. `0,1,2`. (1000 classes)", visible=True)
+        # with gr.Row():
         is_random_checkbox = gr.Checkbox(label="Random shuffle", value=False, elem_id="random_seed_checkbox")
         random_seed_slider = gr.Slider(0, 1000, step=1, label="Random seed", value=1, elem_id="random_seed", visible=False)
-        load_dataset_button = gr.Button("Load Dataset", elem_id="load-dataset-button")
+    
+    if advanced:
+        advanced_block.visible = True
+        basic_block.visible = False
+    else:
+        advanced_block.visible = False
+        basic_block.visible = True
+        
+    # change visibility 
+    advanced_radio.change(fn=lambda x: gr.update(visible=x=="Advanced"), inputs=advanced_radio, outputs=[advanced_block])
+    advanced_radio.change(fn=lambda x: gr.update(visible=x=="Basic"), inputs=advanced_radio, outputs=[basic_block])
+    
     
     def change_filter_options(dataset_name):
         idx = dataset_names.index(dataset_name)
@@ -523,9 +526,12 @@ def make_dataset_images_section(open=False):
     is_random_checkbox.change(fn=change_random_seed, inputs=is_random_checkbox, outputs=random_seed_slider)
     
         
-    def load_dataset_images(dataset_name, num_images=10, 
+    def load_dataset_images(is_advanced, dataset_name, num_images=10, 
                             is_filter=True, filter_by_class_text="0,1,2", 
                             is_random=False, seed=1):
+        if is_advanced == "Basic":
+            gr.Info("Loaded images from Ego-Exo4D")
+            return default_images
         try:
             dataset = load_dataset(dataset_name, trust_remote_code=True)
             key = list(dataset.keys())[0]
@@ -538,7 +544,6 @@ def make_dataset_images_section(open=False):
         
         if is_filter:
             classes = [int(i) for i in filter_by_class_text.split(",")]
-            print(filter_by_class_text, len(classes), classes)
             labels = np.array(dataset['label'])
             unique_labels = np.unique(labels)
             valid_classes = [i for i in classes if i in unique_labels]
@@ -564,15 +569,16 @@ def make_dataset_images_section(open=False):
             else:
                 image_idx = list(range(num_images))
         images = [dataset[i]['image'] for i in image_idx]
+        gr.Info(f"Loaded {len(images)} images from {dataset_name}")
         return images   
     
-    load_dataset_button.click(load_dataset_images, 
-                        inputs=[dataset_dropdown, num_images_slider,
+    load_images_button.click(load_dataset_images, 
+                        inputs=[advanced_radio, dataset_dropdown, num_images_slider,
                                 filter_by_class_checkbox, filter_by_class_text, 
                                 is_random_checkbox, random_seed_slider],
                         outputs=[input_gallery])
     
-    return dataset_dropdown, num_images_slider, random_seed_slider, load_dataset_button
+    return dataset_dropdown, num_images_slider, random_seed_slider, load_images_button
     
 def make_output_images_section():
     gr.Markdown('### Output Images')
@@ -628,8 +634,7 @@ with demo:
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 input_gallery, submit_button, clear_images_button = make_input_images_section()
-                load_images_button, example_gallery, hide_button = make_example_images_section()
-                dataset_dropdown, num_images_slider, random_seed_slider, load_dataset_button = make_dataset_images_section()
+                dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_dataset_images_section()
                 
             with gr.Column(scale=5, min_width=200):
                 output_gallery = make_output_images_section()
@@ -642,7 +647,6 @@ with demo:
                 ] = make_parameters_section()
                 # logging text box
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
-                load_images_button.click(lambda x: default_images, outputs=input_gallery)
 
         clear_images_button.click(lambda x: ([], []), outputs=[input_gallery, output_gallery])
         submit_button.click(
@@ -682,10 +686,7 @@ with demo:
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 input_gallery, submit_button, clear_images_button = make_input_images_section()
-                load_images_button, example_gallery, hide_button = make_example_images_section()
-                dataset_dropdown, num_images_slider, random_seed_slider, load_dataset_button = make_dataset_images_section()
-                example_gallery.visible = False
-                hide_button.visible = False
+                dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_dataset_images_section()
                 
             with gr.Column(scale=5, min_width=200):
                 output_gallery = make_output_images_section()
@@ -703,7 +704,6 @@ with demo:
                     item.visible = False
                 # logging text box
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
-                load_images_button.click(lambda x: (default_images, default_outputs_independent), outputs=[input_gallery, output_gallery])
             
         clear_images_button.click(lambda x: ([], []), outputs=[input_gallery, output_gallery])
         submit_button.click(
@@ -727,11 +727,7 @@ with demo:
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 input_gallery, submit_button, clear_images_button = make_input_images_section()
-                load_images_button, example_gallery, hide_button = make_example_images_section()
-                load_images_button.click(lambda x: default_images, outputs=[input_gallery])
-                example_gallery.visible = False
-                hide_button.visible = False
-                dataset_dropdown, num_images_slider, random_seed_slider, load_dataset_button = make_dataset_images_section()
+                dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_dataset_images_section()
                 num_images_slider.value = 100
                 
             with gr.Column(scale=5, min_width=200):
@@ -791,8 +787,7 @@ with demo:
     with gr.Tab('Video'): 
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
-                video_input_gallery, submit_button, clear_images_button, max_frame_number = make_input_video_section()
-                # load_video_button = make_example_video_section()
+                video_input_gallery, submit_button, clear_video_button, max_frame_number = make_input_video_section()
             with gr.Column(scale=5, min_width=200):
                 video_output_gallery = gr.Video(value=None, label="NCUT Embedding", elem_id="ncut", height="auto", show_share_button=False)
                 [
@@ -808,8 +803,7 @@ with demo:
                 knn_tsne_slider.value = 20
                 # logging text box
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
-                load_images_button.click(lambda x: (default_images, default_outputs), outputs=[input_gallery, output_gallery])
-        clear_images_button.click(lambda x: (None, []), outputs=[video_input_gallery, video_output_gallery])
+        clear_video_button.click(lambda x: (None, None), outputs=[video_input_gallery, video_output_gallery])
         place_holder_false = gr.Checkbox(label="Place holder", value=False, elem_id="place_holder_false")
         place_holder_false.visible = False
         submit_button.click(
@@ -830,12 +824,12 @@ with demo:
         gr.Markdown('---')
         gr.Markdown('![ncut](https://ncut-pytorch.readthedocs.io/en/latest/images/gallery/llama3/llama3_layer_31.jpg)')
     
-    with gr.Tab('Compare'): 
+    with gr.Tab('Compare Models'): 
         def add_one_model(i_model=1):
             with gr.Column(scale=5, min_width=200) as col:
                 gr.Markdown(f'### Output Images')
                 output_gallery = gr.Gallery(value=[], label="NCUT Embedding", show_label=False, elem_id=f"ncut{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto")
-                submit_button = gr.Button("🔴 RUN", elem_id=f"submit_button{i_model}")
+                submit_button = gr.Button("🔴 RUN", elem_id=f"submit_button{i_model}", variant='primary')
                 [
                     model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
                     affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
@@ -863,11 +857,7 @@ with demo:
                 input_gallery, submit_button, clear_images_button = make_input_images_section()
                 clear_images_button.click(lambda x: ([],), outputs=[input_gallery])
                 submit_button.visible = False
-                load_images_button, example_gallery, hide_button = make_example_images_section()
-                example_gallery.visible = False
-                hide_button.visible = False
-                dataset_dropdown, num_images_slider, random_seed_slider, load_dataset_button = make_dataset_images_section(open=True)
-                load_images_button.click(lambda x: default_images, outputs=input_gallery)
+                dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_dataset_images_section(advanced=True)
 
                 
             for i in range(2):
@@ -886,7 +876,7 @@ with demo:
                     with gr.Column(scale=5, min_width=200):
                         add_one_model()
 
-            button = gr.Button("Add Compare", elem_id=f"add_button_{i}", visible=False if i > 0 else True)
+            button = gr.Button("➕ Add Compare", elem_id=f"add_button_{i}", visible=False if i > 0 else True, scale=3)
             buttons.append(button)
             
             if i > 0: