update ncut(legacy)

app.py

@@ -2,8 +2,11 @@
 # %%
 USE_SPACES = True
 
-if USE_SPACES:
-    import spaces
+if USE_SPACES:  # huggingface ZeroGPU
+    try:
+        import spaces
+    except ImportError:
+        USE_SPACES = False  # run on standard GPU
 
 import gradio as gr
 
@@ -14,10 +17,7 @@ import time
 
 import gradio as gr
 
-if USE_SPACES:
-    from backbone import extract_features
-else:
-    from draft_gradio_backbone import extract_features
+from backbone import extract_features
 from ncut_pytorch import NCUT, eigenvector_to_rgb
 
 
@@ -88,9 +88,11 @@ def to_pil_images(images):
         Image.fromarray((image * 255).cpu().numpy().astype(np.uint8)).resize((256, 256), Image.Resampling.NEAREST)
         for image in images
     ]
+    
 
 default_images = ['./images/image_0.jpg', './images/image_1.jpg', './images/image_2.jpg', './images/image_3.jpg', './images/image_5.jpg']
 default_outputs = ['./images/ncut_0.jpg', './images/ncut_1.jpg', './images/ncut_2.jpg', './images/ncut_3.jpg', './images/ncut_5.jpg']
+default_outputs_independent = ['./images/ncut_0_independent.jpg', './images/ncut_1_independent.jpg', './images/ncut_2_independent.jpg', './images/ncut_3_independent.jpg', './images/ncut_5_independent.jpg']
 
 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 = ['./images/ncut_0_small.jpg', './images/ncut_1_small.jpg', './images/ncut_2_small.jpg', './images/ncut_3_small.jpg', './images/ncut_5_small.jpg']
@@ -113,6 +115,7 @@ def ncut_run(
     n_neighbors=500,
     min_dist=0.1,
     sampling_method="fps",
+    old_school_ncut=False,
 ):
     logging_str = ""
     if perplexity >= num_sample_tsne or n_neighbors >= num_sample_tsne:
@@ -134,22 +137,44 @@ def ncut_run(
     # print(f"Feature extraction time (gpu): {time.time() - start:.2f}s")
     logging_str += f"Backbone time: {time.time() - start:.2f}s\n"
     
-    rgb, _logging_str = 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,
-        perplexity=perplexity,
-        n_neighbors=n_neighbors,
-        min_dist=min_dist,
-        sampling_method=sampling_method,
-    )
-    logging_str += _logging_str
-    rgb = dont_use_too_much_green(rgb)
+    if not old_school_ncut:  # joint across all images
+        rgb, _logging_str = 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,
+            perplexity=perplexity,
+            n_neighbors=n_neighbors,
+            min_dist=min_dist,
+            sampling_method=sampling_method,
+        )
+        logging_str += _logging_str
+        rgb = dont_use_too_much_green(rgb)
+    if old_school_ncut:  # individual images
+        logging_str += "Running NCut for each image independently\n"
+        rgb = []
+        for i_image in range(features.shape[0]):
+            feature = features[i_image]
+            _rgb, _logging_str = compute_ncut(
+                feature[None],
+                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,
+                perplexity=perplexity,
+                n_neighbors=n_neighbors,
+                min_dist=min_dist,
+                sampling_method=sampling_method,
+            )
+            logging_str += _logging_str
+            rgb.append(_rgb[0])
     return to_pil_images(rgb), logging_str
 
 def _ncut_run(*args, **kwargs):
@@ -205,6 +230,7 @@ def run_fn(
     n_neighbors=500,
     min_dist=0.1,
     sampling_method="fps",
+    old_school_ncut=False,
 ):
     if images is None:
         gr.Warning("No images selected.")
@@ -228,10 +254,13 @@ def run_fn(
         "n_neighbors": n_neighbors,
         "min_dist": min_dist,
         "sampling_method": sampling_method,
+        "old_school_ncut": old_school_ncut,
     }
     num_images = len(images)
     if num_images > 100:
         return super_duper_long_run(images, **kwargs)
+    if num_images > 50:
+        return longer_run(images, **kwargs)
     if num_images > 10:
         return long_run(images, **kwargs)
     if embedding_method == "UMAP":
@@ -244,73 +273,47 @@ def run_fn(
         return quick_run(images, **kwargs)
     
     return quick_run(images, **kwargs)
-    
-with gr.Blocks() as demo:
 
-    with gr.Row():
-        with gr.Column(scale=5, min_width=200):
-            gr.Markdown('### Input Images')
-            input_gallery = gr.Gallery(value=[], 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')
-        
-            gr.Markdown('### Load from Cloud Dataset 👇')
-            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
-            )
-            
-            with gr.Accordion("➜ Load from dataset", open=True):
-                dataset_names = [
-                    'UCSC-VLAA/Recap-COCO-30K',
-                    'nateraw/pascal-voc-2012',
-                    'johnowhitaker/imagenette2-320',
-                    'jainr3/diffusiondb-pixelart',
-                    'JapanDegitalMaterial/Places_in_Japan',
-                    'Borismile/Anime-dataset',
-                ]
-                dataset_dropdown = gr.Dropdown(dataset_names, label="Dataset name", value="UCSC-VLAA/Recap-COCO-30K", elem_id="dataset")
-                num_images_slider = gr.Slider(1, 200, step=1, label="Number of images", value=9, elem_id="num_images")
-                random_seed_slider = gr.Number(0, label="Random seed", value=42, elem_id="random_seed")
-                load_dataset_button = gr.Button("Load Dataset", elem_id="load-dataset-button")
-
-        with gr.Column(scale=5, min_width=200):
-            gr.Markdown('### Output Images')
-            output_gallery = gr.Gallery(value=[], label="NCUT Embedding", show_label=False, elem_id="ncut", columns=[3], rows=[1], object_fit="contain", height="auto")
-            model_dropdown = gr.Dropdown(["SAM(sam_vit_b)", "MobileSAM", "DiNO(dinov2_vitb14_reg)", "CLIP(openai/clip-vit-base-patch16)", "MAE(vit_base)"], label="Backbone", value="SAM(sam_vit_b)", elem_id="model_name")
-            layer_slider = gr.Slider(0, 11, step=1, label="Backbone: Layer index", value=11, elem_id="layer")
-            node_type_dropdown = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?")
-            num_eig_slider = gr.Slider(1, 1000, step=1, label="NCUT: Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more clusters')
-            affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="NCUT: Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for shaper segmentation")
-            
-            with gr.Accordion("➜ Click to expand: more parameters", open=False):
-                num_sample_ncut_slider = gr.Slider(100, 50000, step=100, label="NCUT: num_sample", value=10000, elem_id="num_sample_ncut", info="Nyström approximation")
-                sampling_method_dropdown = gr.Dropdown(["fps", "random"], label="NCUT: Sampling method", value="fps", elem_id="sampling_method")
-                knn_ncut_slider = gr.Slider(1, 100, step=1, label="NCUT: KNN", value=10, elem_id="knn_ncut", info="Nyström approximation")
-                embedding_method_dropdown = gr.Dropdown(["tsne_3d", "umap_3d", "umap_shpere", "tsne_2d", "umap_2d"], label="Coloring method", value="tsne_3d", elem_id="embedding_method")
-                num_sample_tsne_slider = gr.Slider(100, 1000, step=100, label="t-SNE/UMAP: num_sample", value=300, elem_id="num_sample_tsne", info="Nyström approximation")
-                knn_tsne_slider = gr.Slider(1, 100, step=1, label="t-SNE/UMAP: KNN", value=10, elem_id="knn_tsne", info="Nyström approximation")
-                perplexity_slider = gr.Slider(10, 500, step=10, label="t-SNE: Perplexity", value=150, elem_id="perplexity")
-                n_neighbors_slider = gr.Slider(10, 500, step=10, label="UMAP: n_neighbors", value=150, elem_id="n_neighbors")
-                min_dist_slider = gr.Slider(0.1, 1, step=0.1, label="UMAP: min_dist", value=0.1, elem_id="min_dist")
-
-            # logging text box
-            logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
-            
-    def load_default_images():
-        return default_images, default_outputs
 
-    def empty_input_and_output():
-        return [], []
 
+def make_input_images_section():
+    gr.Markdown('### Input Images')
+    input_gallery = gr.Gallery(value=[], 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')
+    return input_gallery, submit_button, clear_images_button
+
+def make_example_images_section():
+    gr.Markdown('### Load from Cloud Dataset 👇')
+    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
+    )
+    return load_images_button, example_gallery, hide_button
+
+def make_dataset_images_section():
+    with gr.Accordion("➜ Load from dataset", open=True):
+        dataset_names = [
+            'UCSC-VLAA/Recap-COCO-30K',
+            'nateraw/pascal-voc-2012',
+            'johnowhitaker/imagenette2-320',
+            'jainr3/diffusiondb-pixelart',
+            'JapanDegitalMaterial/Places_in_Japan',
+            'Borismile/Anime-dataset',
+        ]
+        dataset_dropdown = gr.Dropdown(dataset_names, label="Dataset name", value="UCSC-VLAA/Recap-COCO-30K", elem_id="dataset")
+        num_images_slider = gr.Slider(1, 200, step=1, label="Number of images", value=9, elem_id="num_images")
+        random_seed_slider = gr.Number(0, label="Random seed", value=42, elem_id="random_seed")
+        load_dataset_button = gr.Button("Load Dataset", elem_id="load-dataset-button")
     def load_dataset_images(dataset_name, num_images=10, random_seed=42):
         from datasets import load_dataset
         try:
-            dataset = load_dataset(dataset_name)['train']
+            dataset = load_dataset(dataset_name)
+            key = list(dataset.keys())[0]
+            dataset = dataset[key]
         except Exception as e:
             gr.Error(f"Error loading dataset {dataset_name}: {e}")
             return None
@@ -319,27 +322,135 @@ with gr.Blocks() as demo:
         image_idx = np.random.RandomState(random_seed).choice(len(dataset), num_images, replace=False)
         image_idx = image_idx.tolist()
         images = [dataset[i]['image'] for i in image_idx]
-        return images    
-
-    
-    load_images_button.click(load_default_images, outputs=[input_gallery, output_gallery])
-    clear_images_button.click(empty_input_and_output, outputs=[input_gallery, output_gallery])
+        return images   
     load_dataset_button.click(load_dataset_images, inputs=[dataset_dropdown, num_images_slider, random_seed_slider], outputs=[input_gallery])
-    submit_button.click(
-        run_fn,
-        inputs=[
-            input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+    return dataset_dropdown, num_images_slider, random_seed_slider, load_dataset_button
+    
+def make_output_images_section():
+    gr.Markdown('### Output Images')
+    output_gallery = gr.Gallery(value=[], label="NCUT Embedding", show_label=False, elem_id="ncut", columns=[3], rows=[1], object_fit="contain", height="auto")
+    return output_gallery
+
+def make_parameters_section():
+    gr.Markdown('### Parameters')
+    model_dropdown = gr.Dropdown(["SAM(sam_vit_b)", "MobileSAM", "DiNO(dinov2_vitb14_reg)", "CLIP(openai/clip-vit-base-patch16)", "MAE(vit_base)"], label="Backbone", value="SAM(sam_vit_b)", elem_id="model_name")
+    layer_slider = gr.Slider(0, 11, step=1, label="Backbone: Layer index", value=11, elem_id="layer")
+    node_type_dropdown = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?")
+    num_eig_slider = gr.Slider(1, 1000, step=1, label="NCUT: Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more clusters')
+    affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="NCUT: Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for shaper segmentation")
+
+    with gr.Accordion("➜ Click to expand: more parameters", open=False):
+        num_sample_ncut_slider = gr.Slider(100, 50000, step=100, label="NCUT: num_sample", value=10000, elem_id="num_sample_ncut", info="Nyström approximation")
+        sampling_method_dropdown = gr.Dropdown(["fps", "random"], label="NCUT: Sampling method", value="fps", elem_id="sampling_method")
+        knn_ncut_slider = gr.Slider(1, 100, step=1, label="NCUT: KNN", value=10, elem_id="knn_ncut", info="Nyström approximation")
+        embedding_method_dropdown = gr.Dropdown(["tsne_3d", "umap_3d", "umap_shpere", "tsne_2d", "umap_2d"], label="Coloring method", value="tsne_3d", elem_id="embedding_method")
+        num_sample_tsne_slider = gr.Slider(100, 1000, step=100, label="t-SNE/UMAP: num_sample", value=300, elem_id="num_sample_tsne", info="Nyström approximation")
+        knn_tsne_slider = gr.Slider(1, 100, step=1, label="t-SNE/UMAP: KNN", value=10, elem_id="knn_tsne", info="Nyström approximation")
+        perplexity_slider = gr.Slider(10, 500, step=10, label="t-SNE: Perplexity", value=150, elem_id="perplexity")
+        n_neighbors_slider = gr.Slider(10, 500, step=10, label="UMAP: n_neighbors", value=150, elem_id="n_neighbors")
+        min_dist_slider = gr.Slider(0.1, 1, step=0.1, label="UMAP: min_dist", value=0.1, elem_id="min_dist")
+    return [model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
             affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
             embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
-            perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
-        ],
-        outputs=[output_gallery, logging_text]
-    )
+            perplexity_slider, n_neighbors_slider, min_dist_slider, 
+            sampling_method_dropdown]
 
+with gr.Blocks() as demo:
 
-if USE_SPACES:
-    demo.launch()
-else:
-    demo.launch(share=True)
+    with gr.Tab('AlignedCut (Images)'):
+
+        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()
+                
+            with gr.Column(scale=5, min_width=200):
+                output_gallery = make_output_images_section()
+                [
+                    model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
+                    affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
+                    embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                    perplexity_slider, n_neighbors_slider, min_dist_slider, 
+                    sampling_method_dropdown
+                ] = 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, default_outputs), outputs=[input_gallery, output_gallery])
+
+        clear_images_button.click(lambda x: ([], []), outputs=[input_gallery, output_gallery])
+        submit_button.click(
+            run_fn,
+            inputs=[
+                input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
+                embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
+            ],
+            outputs=[output_gallery, logging_text]
+        )
+        
+    with gr.Tab('NCut (Legacy)'): 
+        gr.Markdown('#### Ncut, not aligned, no Nyström approximation')
+        gr.Markdown('1. Each image is solved independently, _color is not aligned across images_')
+        gr.Markdown('2. No Nyström approximation')
+        
+        gr.Markdown('### NCut (Legacy) vs. AlignedCut:')
+        with gr.Row():
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('#### Pros')
+                gr.Markdown('- Easy Solution. Use less eigenvectors.')
+                gr.Markdown('- Exact solution. No Nyström approximation.')
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('#### Cons')
+                gr.Markdown('- Not aligned. Distance is not preserved across images. No pseudo-labeling or correspondence.')
+                gr.Markdown('- Poor complexity scaling. Unable to handle large number of pixels.')
+        gr.Markdown('---')
+        with gr.Row():
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown(' ')
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('_color is not aligned across images_ 👇')
+
+
+        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
+                
+            with gr.Column(scale=5, min_width=200):
+                output_gallery = make_output_images_section()
+                [
+                    model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
+                    affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
+                    embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                    perplexity_slider, n_neighbors_slider, min_dist_slider, 
+                    sampling_method_dropdown
+                ] = make_parameters_section()
+                old_school_ncut_checkbox = gr.Checkbox(label="Old school NCut", value=True, elem_id="old_school_ncut")
+                invisible_list = [old_school_ncut_checkbox, num_sample_ncut_slider, knn_ncut_slider,
+                                    num_sample_tsne_slider, knn_tsne_slider, sampling_method_dropdown]
+                for item in invisible_list:
+                    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(
+            run_fn,
+            inputs=[
+                input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
+                embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown, 
+                old_school_ncut_checkbox
+            ],
+            outputs=[output_gallery, logging_text]
+        )
+demo.launch(share=True)
 
 # %%