update cluster fg bg

app.py

@@ -308,7 +308,79 @@ def blend_image_with_heatmap(image, heatmap, opacity1=0.5, opacity2=0.5):
     blended = (1 - opacity1) * image + opacity2 * heatmap
     return blended.astype(np.uint8)
 
-def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=False):
+
+def segment_fg_bg(images):
+
+    images = F.interpolate(images, (224, 224), mode="bilinear")
+
+    # model = load_alignedthreemodel()
+    model = load_model("CLIP(ViT-B-16/openai)")
+    from ncut_pytorch.backbone import resample_position_embeddings
+    pos_embed = model.model.visual.positional_embedding
+    pos_embed = resample_position_embeddings(pos_embed, 14, 14)
+    model.model.visual.positional_embedding = torch.nn.Parameter(pos_embed)
+    
+    batch_size = 4
+    chunk_idxs = torch.split(torch.arange(images.shape[0]), batch_size)
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model.to(device)
+    means = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
+    stds = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
+    
+    fg_acts, bg_acts = [], []
+    for chunk_idx in chunk_idxs:
+        with torch.no_grad():
+            input_images = images[chunk_idx].to(device)
+            # transform the input images
+            input_images = (input_images - means) / stds
+            # output = model(input_images)[:, 5]
+            output = model(input_images)['attn'][6]
+            fg_act = output[:, 6, 6].mean(0)
+            bg_act = output[:, 0, 0].mean(0)
+            fg_acts.append(fg_act)
+            bg_acts.append(bg_act)
+    fg_act = torch.stack(fg_acts, dim=0).mean(0)
+    bg_act = torch.stack(bg_acts, dim=0).mean(0)
+    fg_act = F.normalize(fg_act, dim=-1)
+    bg_act = F.normalize(bg_act, dim=-1)
+    
+    # ref_image = default_images[0]
+    # image = Image.open(ref_image).convert("RGB").resize((224, 224), Image.Resampling.BILINEAR)
+    # image = torch.tensor(np.array(image)).permute(2, 0, 1).float().to(device)
+    # image = (image / 255.0 - means) / stds
+    # output = model(image)['attn'][6][0]
+    # # print(output.shape)
+    # # bg on the center
+    # fg_act = output[5, 5]
+    # # bg on the bottom left
+    # bg_act = output[0, 0]
+    # fg_act = F.normalize(fg_act, dim=-1)
+    # bg_act = F.normalize(bg_act, dim=-1)
+    
+    # print(images.mean(), images.std())
+    
+    fg_act, bg_act = fg_act.to(device), bg_act.to(device)
+    chunk_idxs = torch.split(torch.arange(images.shape[0]), batch_size)
+    heatmap_fgs, heatmap_bgs = [], []
+    for chunk_idx in chunk_idxs:
+        with torch.no_grad():
+            input_images = images[chunk_idx].to(device)
+            # transform the input images
+            input_images = (input_images - means) / stds
+            # output = model(input_images)[:, 5]
+            output = model(input_images)['attn'][6]
+            output = F.normalize(output, dim=-1)
+            heatmap_fg = output @ fg_act[:, None]
+            heatmap_bg = output @ bg_act[:, None]
+            heatmap_fgs.append(heatmap_fg.cpu())
+            heatmap_bgs.append(heatmap_bg.cpu())
+    heatmap_fg = torch.cat(heatmap_fgs, dim=0)
+    heatmap_bg = torch.cat(heatmap_bgs, dim=0)
+    return heatmap_fg, heatmap_bg
+
+
+def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=False, clusters=50, eig_idx=None, title='cluster'):
     progress = gr.Progress()
     progress(progess_start, desc="Finding Clusters by FPS")
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -318,10 +390,13 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
     
     # gr.Info("Finding Clusters by FPS, no magnitude filtering")
     top_p_idx = torch.arange(eigvecs.shape[0])
+    if eig_idx is not None:
+        top_p_idx = eig_idx
     # gr.Info("Finding Clusters by FPS, with magnitude filtering")
     # p = 0.8
     # top_p_idx = magnitude.argsort(descending=True)[:int(p * magnitude.shape[0])]
     
+    
     ret_magnitude = magnitude.reshape(-1, h, w)
     
     
@@ -338,7 +413,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
     right = F.normalize(right, dim=-1)
     heatmap = left @ right.T
     heatmap = F.normalize(heatmap, dim=-1)
-    num_samples = 50 if not advanced else 100
+    num_samples = clusters + 20
     if num_samples > fps_idx.shape[0]:
         num_samples = fps_idx.shape[0]
     r2_fps_idx = farthest_point_sampling(heatmap, num_samples)
@@ -398,10 +473,10 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
         
     fig_images = []
     i_cluster = 0
-    num_plots = 10 if not advanced else 20
+    num_plots = clusters // 5
     plot_step_float = (1.0 - progess_start) / num_plots
     for i_fig in range(num_plots):    
-        progress(progess_start + i_fig * plot_step_float, desc="Plotting Clusters")
+        progress(progess_start + i_fig * plot_step_float, desc=f"Plotting {title}")
         if not advanced:
             fig, axs = plt.subplots(3, 5, figsize=(15, 9))
         if advanced:
@@ -421,7 +496,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
                 _heatmap = blend_image_with_heatmap(images[image_idx], heatmap[i])
                 axs[i, j].imshow(_heatmap)
                 if i == 0:
-                    axs[i, j].set_title(f"cluster {i_cluster+1}", fontsize=24)
+                    axs[i, j].set_title(f"{title} {i_cluster+1}", fontsize=24)
                     i_cluster += 1
         plt.tight_layout(h_pad=0.5, w_pad=0.3)
 
@@ -440,6 +515,39 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
 
     return fig_images, ret_magnitude
 
+def make_cluster_plot_advanced(eigvecs, images, h=64, w=64):
+    heatmap_fg, heatmap_bg = segment_fg_bg(images.clone())
+    heatmap_bg = rearrange(heatmap_bg, 'b h w c -> b c h w')
+    heatmap_fg = rearrange(heatmap_fg, 'b h w c -> b c h w')
+    heatmap_fg = F.interpolate(heatmap_fg, (h, w), mode="bilinear")
+    heatmap_bg = F.interpolate(heatmap_bg, (h, w), mode="bilinear")
+    heatmap_fg = heatmap_fg.flatten()
+    heatmap_bg = heatmap_bg.flatten()
+    
+    fg_minus_bg = heatmap_fg - heatmap_bg
+    fg_mask = fg_minus_bg > fg_minus_bg.quantile(0.8)
+    bg_mask = fg_minus_bg < fg_minus_bg.quantile(0.2)
+
+    # fg_mask = heatmap_fg > heatmap_fg.quantile(0.8)
+    # bg_mask = heatmap_bg > heatmap_bg.quantile(0.8)
+    other_mask = ~(fg_mask | bg_mask)
+    
+    fg_idx = torch.arange(heatmap_fg.shape[0])[fg_mask]
+    bg_idx = torch.arange(heatmap_bg.shape[0])[bg_mask]
+    other_idx = torch.arange(heatmap_fg.shape[0])[other_mask]
+    
+    fg_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=100, eig_idx=fg_idx, title="fg")
+    bg_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=100, eig_idx=bg_idx, title="bg")
+    other_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=100, eig_idx=other_idx, title="other")
+    
+    cluster_images = fg_images + bg_images + other_images
+    
+    magitude = torch.norm(eigvecs, dim=-1)
+    magitude = magitude.reshape(-1, h, w)
+    
+    # magitude = fg_minus_bg.reshape(-1, h, w)  #TODO
+    
+    return cluster_images, magitude
     
 def ncut_run(
     model,
@@ -601,7 +709,7 @@ def ncut_run(
                 if torch.cuda.is_available():
                     images = images.cuda()
                 _images = reverse_transform_image(images, stablediffusion="stable" in model_name.lower())
-                cluster_images, eig_magnitude = make_cluster_plot(eigvecs, _images, h=h, w=w, progess_start=progress_start, advanced=advanced)
+                cluster_images, eig_magnitude = make_cluster_plot_advanced(eigvecs, _images, h=h, w=w)
                 logging_str += f"Recursion #{i+1} plot time: {time.time() - start:.2f}s\n"    
                 
                 norm_images = []
@@ -716,7 +824,10 @@ def ncut_run(
             images = images.cuda()
         _images = reverse_transform_image(images, stablediffusion="stable" in model_name.lower())
         advanced = kwargs.get("advanced", False)
-        cluster_images, eig_magnitude = make_cluster_plot(eigvecs, _images, h=h, w=w, progess_start=progress_start, advanced=advanced)
+        if advanced:
+            cluster_images, eig_magnitude = make_cluster_plot_advanced(eigvecs, _images, h=h, w=w)
+        else:
+            cluster_images, eig_magnitude = make_cluster_plot(eigvecs, _images, h=h, w=w, progess_start=progress_start, advanced=False)
         logging_str += f"plot time: {time.time() - start:.2f}s\n"
     
     norm_images = None
@@ -736,33 +847,33 @@ def ncut_run(
         logging_str += "Eigenvector Magnitude\n"
         logging_str += f"Min: {vmin:.2f}, Max: {vmax:.2f}\n"
         gr.Info(f"Eigenvector Magnitude:</br> Min: {vmin:.2f}, Max: {vmax:.2f}", duration=10)
-
+    
     return to_pil_images(rgb), cluster_images, norm_images, logging_str
 
 
 
 def _ncut_run(*args, **kwargs):
     n_ret = kwargs.pop("n_ret", 1)
-    try:
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+    # try:
+    #     if torch.cuda.is_available():
+    #         torch.cuda.empty_cache()
             
-        ret = ncut_run(*args, **kwargs)
+    #     ret = ncut_run(*args, **kwargs)
         
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+    #     if torch.cuda.is_available():
+    #         torch.cuda.empty_cache()
         
-        ret = list(ret)[:n_ret] + [ret[-1]]
-        return ret
-    except Exception as e:
-        gr.Error(str(e))
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        return *(None for _ in range(n_ret)), "Error: " + str(e)
-
-    # ret = ncut_run(*args, **kwargs)
-    # ret = list(ret)[:n_ret] + [ret[-1]]
-    # return ret
+    #     ret = list(ret)[:n_ret] + [ret[-1]]
+    #     return ret
+    # except Exception as e:
+    #     gr.Error(str(e))
+    #     if torch.cuda.is_available():
+    #         torch.cuda.empty_cache()
+    #     return *(None for _ in range(n_ret)), "Error: " + str(e)
+
+    ret = ncut_run(*args, **kwargs)
+    ret = list(ret)[:n_ret] + [ret[-1]]
+    return ret
 
 if USE_HUGGINGFACE_ZEROGPU:
     @spaces.GPU(duration=30)
@@ -1186,7 +1297,7 @@ def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_
                 images += [Image.open(new_image) for new_image in new_images]
             if isinstance(new_images, str):
                 images.append(Image.open(new_images))
-            new_images = None
+            gr.Info(f"Total images: {len(images)}")
             return images
         upload_button.upload(convert_to_pil_and_append, inputs=[input_gallery, upload_button], outputs=[input_gallery])
         
@@ -1402,6 +1513,7 @@ def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_
         if existing_images is None:
             existing_images = []
         existing_images += new_images
+        gr.Info(f"Total images: {len(existing_images)}")
         return existing_images
     
     load_images_button.click(load_and_append, 
@@ -1416,165 +1528,6 @@ def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_
     
 
 
-# def make_input_images_section(rows=1, cols=3, height="auto"):
-#     gr.Markdown('### Input Images')
-#     input_gallery = gr.Gallery(value=None, label="Select images", show_label=True, elem_id="images", columns=[cols], rows=[rows], object_fit="contain", height=height, type="pil", show_share_button=False)
-#     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_dataset_images_section(advanced=False, is_random=False):
-
-#     gr.Markdown('### Load Datasets')
-#     load_images_button = gr.Button("🔴 Load Images", elem_id="load-images-button", variant='primary')
-#     advanced_radio = gr.Radio(["Basic", "Advanced"], label="Datasets", value="Advanced" if advanced else "Basic", elem_id="advanced-radio", show_label=True)
-#     with gr.Column() as basic_block:
-#         example_gallery = gr.Gallery(value=example_items, label="Example Images", show_label=True, 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_categories = list(DATASETS.keys())
-#         defualt_cat = dataset_categories[0]
-#         def get_choices(cat):
-#             return [tup[0] for tup in DATASETS[cat]]
-#         defualt_choices = get_choices(defualt_cat)
-#         with gr.Row():
-#             dataset_radio = gr.Radio(dataset_categories, label="Dataset Category", value=defualt_cat, elem_id="dataset-radio", show_label=True, min_width=600)
-#             # dataset_dropdown = gr.Dropdown(dataset_names, label="Dataset name", value="mrm8488/ImageNet1K-val", elem_id="dataset", min_width=300)
-#             dataset_dropdown = gr.Dropdown(defualt_choices, label="Dataset name", value=defualt_choices[0], elem_id="dataset", min_width=400)
-#             dataset_radio.change(fn=lambda x: gr.update(choices=get_choices(x), value=get_choices(x)[0]), inputs=dataset_radio, outputs=dataset_dropdown)
-#             # num_images_slider = gr.Number(10, label="Number of images", elem_id="num_images")
-#             num_images_slider = gr.Slider(1, 1000, step=1, label="Number of images", value=10, elem_id="num_images", min_width=200)
-#             if not is_random:
-#                 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)
-#                 # 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)
-#                 is_random_checkbox = gr.Checkbox(label="Random shuffle", value=True, elem_id="random_seed_checkbox")
-#                 random_seed_slider = gr.Slider(0, 1000, step=1, label="Random seed", value=1, elem_id="random_seed", visible=True)
-#             if is_random:
-#                 filter_by_class_checkbox = gr.Checkbox(label="Filter by class", value=False, 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=False)
-#                 is_random_checkbox = gr.Checkbox(label="Random shuffle", value=True, elem_id="random_seed_checkbox")
-#                 random_seed_slider = gr.Slider(0, 1000, step=1, label="Random seed", value=42, elem_id="random_seed", visible=True)
-                
-    
-#     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 find_num_classes(dataset_name):
-#         num_classes = None
-#         for cat, datasets in DATASETS.items():
-#             datasets = [tup[0] for tup in datasets]
-#             if dataset_name in datasets:
-#                 num_classes = DATASETS[cat][datasets.index(dataset_name)][1]
-#                 break
-#         return num_classes
-    
-#     def change_filter_options(dataset_name):
-#         num_classes = find_num_classes(dataset_name)
-#         if num_classes is None:
-#             return (gr.Checkbox(label="Filter by class", value=False, elem_id="filter_by_class_checkbox", visible=False),
-#                     gr.Textbox(label="Class to select", value="0,1,2", elem_id="filter_by_class_text", info="e.g. `0,1,2`. This dataset has no class label", visible=False))
-#         return (gr.Checkbox(label="Filter by class", value=True, elem_id="filter_by_class_checkbox", visible=True),
-#                 gr.Textbox(label="Class to select", value="0,1,2", elem_id="filter_by_class_text", info=f"e.g. `0,1,2`. ({num_classes} classes)", visible=True))
-#     dataset_dropdown.change(fn=change_filter_options, inputs=dataset_dropdown, outputs=[filter_by_class_checkbox, filter_by_class_text])
-    
-#     def change_filter_by_class(is_filter, dataset_name):
-#         num_classes = find_num_classes(dataset_name)
-#         return gr.Textbox(label="Class to select", value="0,1,2", elem_id="filter_by_class_text", info=f"e.g. `0,1,2`. ({num_classes} classes)", visible=is_filter)
-#     filter_by_class_checkbox.change(fn=change_filter_by_class, inputs=[filter_by_class_checkbox, dataset_dropdown], outputs=filter_by_class_text)
-    
-#     def change_random_seed(is_random):
-#         return gr.Slider(0, 1000, step=1, label="Random seed", value=1, elem_id="random_seed", visible=is_random)
-#     is_random_checkbox.change(fn=change_random_seed, inputs=is_random_checkbox, outputs=random_seed_slider)
-    
-        
-#     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):
-#         progress = gr.Progress()
-#         progress(0, desc="Loading Images")
-#         if is_advanced == "Basic":
-#             gr.Info("Loaded images from Ego-Exo4D")
-#             return default_images
-#         try:
-#             progress(0.5, desc="Downloading Dataset")
-#             dataset = load_dataset(dataset_name, trust_remote_code=True)
-#             key = list(dataset.keys())[0]
-#             dataset = dataset[key]
-#         except Exception as e:
-#             gr.Error(f"Error loading dataset {dataset_name}: {e}")
-#             return None
-#         if num_images > len(dataset):
-#             num_images = len(dataset)
-        
-#         if is_filter:
-#             progress(0.8, desc="Filtering Images")
-#             classes = [int(i) for i in filter_by_class_text.split(",")]
-#             labels = np.array(dataset['label'])
-#             unique_labels = np.unique(labels)
-#             valid_classes = [i for i in classes if i in unique_labels]
-#             invalid_classes = [i for i in classes if i not in unique_labels]
-#             if len(invalid_classes) > 0:
-#                 gr.Warning(f"Classes {invalid_classes} not found in the dataset.")
-#             if len(valid_classes) == 0:
-#                 gr.Error(f"Classes {classes} not found in the dataset.")
-#                 return None
-#             # shuffle each class
-#             chunk_size = num_images // len(valid_classes)
-#             image_idx = []
-#             for i in valid_classes:
-#                 idx = np.where(labels == i)[0]
-#                 if is_random:
-#                     idx = np.random.RandomState(seed).choice(idx, chunk_size, replace=False)
-#                 else:
-#                     idx = idx[:chunk_size]
-#                 image_idx.extend(idx.tolist())
-#         if not is_filter:
-#             if is_random:
-#                 image_idx = np.random.RandomState(seed).choice(len(dataset), num_images, replace=False).tolist()
-#             else:
-#                 image_idx = list(range(num_images))
-#         key = 'image' if 'image' in dataset[0] else list(dataset[0].keys())[0]
-#         images = [dataset[i][key] for i in image_idx]
-#         gr.Info(f"Loaded {len(images)} images from {dataset_name}")
-#         del dataset
-        
-#         if dataset_name in CENTER_CROP_DATASETS:
-#             def center_crop_image(img):
-#                 # image: PIL image
-#                 w, h = img.size
-#                 min_hw = min(h, w)
-#                 # center crop
-#                 left = (w - min_hw) // 2
-#                 top = (h - min_hw) // 2
-#                 right = left + min_hw
-#                 bottom = top + min_hw
-#                 img = img.crop((left, top, right, bottom))
-#                 return img
-#             images = [center_crop_image(image) for image in images]
-        
-#         return images   
-    
-#     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_images_button
-    
-
 # def random_rotate_rgb_gallery(images):
 #     if images is None or len(images) == 0:
 #         gr.Warning("No images selected.")
@@ -1969,19 +1922,19 @@ with demo:
                 l1_gallery = gr.Gallery(format='png', value=[], label="Recursion #1", show_label=True, elem_id="ncut_l1", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
                 add_output_images_buttons(l1_gallery)
                 l1_norm_gallery = gr.Gallery(value=[], label="Recursion #1 Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
-                l1_cluster_gallery = gr.Gallery(value=[], label="Recursion #1 Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=600, show_share_button=True, preview=True, interactive=False)
+                l1_cluster_gallery = gr.Gallery(value=[], label="Recursion #1 Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=500, show_share_button=True, preview=True, interactive=False)
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #2)')
                 l2_gallery = gr.Gallery(format='png', value=[], label="Recursion #2", show_label=True, elem_id="ncut_l2", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
                 add_output_images_buttons(l2_gallery)
                 l2_norm_gallery = gr.Gallery(value=[], label="Recursion #2 Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
-                l2_cluster_gallery = gr.Gallery(value=[], label="Recursion #2 Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=600, show_share_button=True, preview=True, interactive=False)
+                l2_cluster_gallery = gr.Gallery(value=[], label="Recursion #2 Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=500, show_share_button=True, preview=True, interactive=False)
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #3)')
                 l3_gallery = gr.Gallery(format='png', value=[], label="Recursion #3", show_label=True, elem_id="ncut_l3", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
                 add_output_images_buttons(l3_gallery)
                 l3_norm_gallery = gr.Gallery(value=[], label="Recursion #3 Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
-                l3_cluster_gallery = gr.Gallery(value=[], label="Recursion #3 Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=600, show_share_button=True, preview=True, interactive=False)
+                l3_cluster_gallery = gr.Gallery(value=[], label="Recursion #3 Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=500, show_share_button=True, preview=True, interactive=False)
         
         with gr.Row():        
             with gr.Column(scale=5, min_width=200):
@@ -2352,7 +2305,7 @@ with demo:
                 submit_button = gr.Button("🔴 RUN", elem_id=f"submit_button{i_model}", variant='primary')
                 add_output_images_buttons(output_gallery)
                 norm_gallery = gr.Gallery(value=[], label="Eigenvector Magnitude", show_label=True, elem_id=f"eig_norm{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
-                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id=f"clusters{i_model}", columns=[2], rows=[4], object_fit="contain", height=600, show_share_button=True, preview=True, interactive=False)
+                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id=f"clusters{i_model}", columns=[2], rows=[4], object_fit="contain", height=500, show_share_button=True, preview=True, interactive=False)
                 [
                     model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
                     affinity_focal_gamma_slider, num_sample_ncut_slider, ncut_knn_slider, ncut_indirect_connection, ncut_make_orthogonal,