add kway

app.py

@@ -927,6 +927,44 @@ def ncut_run(
         
         return to_pil_images(rgb_all), logging_str
     
+    kway = kwargs.get("kway", False)
+    if kway:
+        only_eigvecs = True
+        rgb, _logging_str, eigvecs = compute_ncut(
+            features,
+            num_eig=num_eig,
+            num_sample_ncut=num_sample_ncut,
+            affinity_focal_gamma=affinity_focal_gamma,
+            knn_ncut=knn_ncut,
+            knn_tsne=knn_tsne,
+            num_sample_tsne=num_sample_tsne,
+            embedding_method=embedding_method,
+            embedding_metric=embedding_metric,
+            perplexity=perplexity,
+            n_neighbors=n_neighbors,
+            min_dist=min_dist,
+            sampling_method=sampling_method,
+            indirect_connection=indirect_connection,
+            make_orthogonal=make_orthogonal,
+            metric=ncut_metric,
+            only_eigvecs=only_eigvecs,
+        )
+        from ncut_pytorch import kway_ncut
+        kway_onehot = kway_ncut(eigvecs)  # [N, K]
+        kway_indices = kway_onehot.argmax(dim=-1)  # [N]
+        kway_indices = kway_indices.cpu().numpy()
+        if kway_indices.max() > 10:
+            cm = plt.colormaps['tab20']
+            rgb = cm(kway_indices / 20)
+        else:
+            cm = plt.colormaps['tab10']
+            rgb = cm(kway_indices / 10)
+        if kway_indices.max() > 20:
+            gr.Error("Too many clusters for kway_ncut")
+        rgb = rgb[:, :3]
+        rgb = rgb.reshape(*features.shape[:-1], 3)
+        return to_pil_images(rgb), logging_str
+        
     # ailgnedcut
     if not directed:
         only_eigvecs = kwargs.get("only_eigvecs", False)
@@ -1318,6 +1356,7 @@ def run_fn(
     return_eigvec_and_rgb=False,
     normalize_eigvec_return=False,
     separate_fg_bg=False,
+    kway=False,
 ):
     # print(node_type2, head_index_text, make_symmetric)
     progress=gr.Progress()
@@ -1463,6 +1502,7 @@ def run_fn(
         "return_eigvec_and_rgb": return_eigvec_and_rgb,
         "normalize_eigvec_return": normalize_eigvec_return,
         "separate_fg_bg": separate_fg_bg,
+        "kway": kway,
     }
     # print(kwargs)
     
@@ -4348,6 +4388,44 @@ with demo:
             outputs=[output_gallery, logging_text],
         )
         
+        
+    with gr.Tab('K-way'):
+
+        with gr.Row():
+            with gr.Column(scale=5, min_width=200):
+                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section()
+                num_images_slider.value = 30
+                logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information", autofocus=False, autoscroll=False)
+                
+            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)
+                [
+                    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, 
+                    embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                    perplexity_slider, n_neighbors_slider, min_dist_slider, 
+                    sampling_method_dropdown, ncut_metric_dropdown, positive_prompt, negative_prompt
+                ] = make_parameters_section()
+                num_eig_slider.value = 6
+                num_eig_slider.maximum = 20
+
+        false_placeholder = gr.Checkbox(label="False", value=False, elem_id="false_placeholder", visible=False)
+        no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
+        
+        submit_button.click(
+            partial(run_fn, n_ret=1, plot_clusters=False, kway=True),
+            inputs=[
+                input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                positive_prompt, negative_prompt,
+                false_placeholder, no_prompt, no_prompt, no_prompt,
+                affinity_focal_gamma_slider, num_sample_ncut_slider, ncut_knn_slider, ncut_indirect_connection, ncut_make_orthogonal, 
+                embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown, ncut_metric_dropdown,
+                *[false_placeholder]*12,
+            ],
+            outputs=[output_gallery, logging_text],
+        )
                         
     with gr.Tab('Sub-cluster (dev)', visible=False) as sub_cluster_tab:
         with gr.Row():