update plotting fg

app.py

@@ -312,7 +312,8 @@ def blend_image_with_heatmap(image, heatmap, opacity1=0.5, opacity2=0.5):
     blended = (1 - opacity1) * image + opacity2 * heatmap
     return blended.astype(np.uint8)
 
-
+# preload the model
+load_model("CLIP(ViT-B-16/openai)")
 def segment_fg_bg(images):
 
     images = F.interpolate(images, (224, 224), mode="bilinear")
@@ -459,7 +460,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
         # fps_heatmaps[idx.item()] = heatmap.cpu()
         fps_heatmaps[idx.item()] = heatmap[mask_sort_idx[:6]].cpu()
         top3_image_idx[idx.item()] = mask_sort_idx[:3]
-        top10_image_idx[idx.item()] = mask_sort_idx[:6]
+        top10_image_idx[idx.item()] = mask_sort_idx[:5]
     # do the sorting
     _sort_idx = torch.tensor(sort_values).argsort(descending=True)
     fps_idx = fps_idx[_sort_idx]
@@ -486,7 +487,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
         if not advanced:
             fig, axs = plt.subplots(3, 5, figsize=(15, 9))
         if advanced:
-            fig, axs = plt.subplots(6, 5, figsize=(15, 18))
+            fig, axs = plt.subplots(5, 5, figsize=(15, 15))
         for ax in axs.flatten():
             ax.axis("off")
         for j, idx in enumerate(fps_idx[i_fig*5:i_fig*5+5]):
@@ -521,7 +522,7 @@ 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):
+def make_cluster_plot_advanced(eigvecs, images, h=64, w=64, num_fg=100, num_bg=10, num_other=0, small_title=True):
     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')
@@ -542,9 +543,15 @@ def make_cluster_plot_advanced(eigvecs, images, h=64, w=64):
     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")
+    fg_images = []
+    if num_fg > 0:
+        fg_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=num_fg, eig_idx=fg_idx, title="fg" if small_title else "cluster")
+    bg_images = []
+    if num_bg > 0:
+        bg_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=num_bg, eig_idx=bg_idx, title="bg" if small_title else "cluster")
+    other_images = []
+    if num_other > 0:
+        other_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=num_other, eig_idx=other_idx, title="other" if small_title else "cluster")
     
     cluster_images = fg_images + bg_images + other_images
     
@@ -833,7 +840,7 @@ def ncut_run(
         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)
+            cluster_images, eig_magnitude = make_cluster_plot_advanced(eigvecs, _images, h=h, w=w, num_fg=20, num_bg=0, num_other=0, small_title=False)
         logging_str += f"plot time: {time.time() - start:.2f}s\n"
     
     norm_images = None
@@ -1859,7 +1866,7 @@ with demo:
                 
             with gr.Column(scale=5, min_width=200):
                 output_gallery = make_output_images_section()
-                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[2], 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=[2], rows=[2], object_fit="contain", height="auto", show_share_button=True, preview=True, interactive=True)
                 [
                     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,