remove mask size filter

app.py

@@ -287,7 +287,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 = 80 if not advanced else 130
+    num_samples = 50 if not advanced else 100
     if num_samples > fps_idx.shape[0]:
         num_samples = fps_idx.shape[0]
     r2_fps_idx = farthest_point_sampling(heatmap, num_samples)
@@ -305,6 +305,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
     fps_heatmaps = {}
     sort_values = []
     top3_image_idx = {}
+    top10_image_idx = {}
     for _, idx in enumerate(fps_idx):
         heatmap = F.cosine_similarity(eigvecs, eigvecs[idx][None], dim=-1)
         
@@ -314,7 +315,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
         #         tensor = tensor[torch.randperm(tensor.shape[0])[:max_size]]
         #     return tensor.quantile(p)
         # top_p = top_percentile(heatmap, p=0.5)
-        top_p = 0.5
+        top_p = 0.8
         
         heatmap = heatmap.reshape(-1, h, w)
         mask = (heatmap > top_p).float()
@@ -324,8 +325,9 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
         mask = mask[mask_sort_idx[:3]] 
         sort_values.append(mask.mean().item())
         # fps_heatmaps[idx.item()] = heatmap.cpu()
-        fps_heatmaps[idx.item()] = heatmap[mask_sort_idx[:3]].cpu()
+        fps_heatmaps[idx.item()] = heatmap[mask_sort_idx[:10]].cpu()
         top3_image_idx[idx.item()] = mask_sort_idx[:3]
+        top10_image_idx[idx.item()] = mask_sort_idx[:10]
     # do the sorting
     _sort_idx = torch.tensor(sort_values).argsort(descending=True)
     fps_idx = fps_idx[_sort_idx]
@@ -342,13 +344,17 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
         # shuffle the fps_idx
         fps_idx = fps_idx[torch.randperm(fps_idx.shape[0])]
     
+        
     fig_images = []
     i_cluster = 0
     num_plots = 10 if not advanced else 20
     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")
-        fig, axs = plt.subplots(3, 5, figsize=(15, 9))
+        if not advanced:
+            fig, axs = plt.subplots(3, 5, figsize=(15, 9))
+        if advanced:
+            fig, axs = plt.subplots(6, 5, figsize=(15, 18))
         for ax in axs.flatten():
             ax.axis("off")
         for j, idx in enumerate(fps_idx[i_fig*5:i_fig*5+5]):
@@ -358,7 +364,8 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
             size = (images.shape[1], images.shape[2])
             heatmap = apply_reds_colormap(heatmap, size)
             # for i, image_idx in enumerate(sorted_image_idxs[:3]):
-            for i, image_idx in enumerate(top3_image_idx[idx.item()]):
+            image_idxs = top3_image_idx[idx.item()] if not advanced else top10_image_idx[idx.item()]
+            for i, image_idx in enumerate(image_idxs):
                 # _heatmap = blend_image_with_heatmap(images[image_idx], heatmap[image_idx])
                 _heatmap = blend_image_with_heatmap(images[image_idx], heatmap[i])
                 axs[i, j].imshow(_heatmap)
@@ -378,10 +385,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
 
         fig_images.append(img)
         plt.close()
-        
-        # plt.imshow(img)
-        # plt.axis("off")
-        # plt.show()
+
     return fig_images, ret_magnitude
 
     
@@ -647,26 +651,26 @@ def ncut_run(
 
 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)
@@ -1415,7 +1419,7 @@ with demo:
             with gr.Column(scale=5, min_width=200):
                 output_gallery = make_output_images_section()
                 norm_gallery = gr.Gallery(value=[], label="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)
-                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[5], rows=[2], object_fit="contain", height="auto", show_share_button=True, preview=True, interactive=False)
+                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=600, 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,