update dataset

app.py

@@ -1,4 +1,10 @@
-import spaces
+# Author: Huzheng Yang
+# %%
+USE_SPACES = False
+
+if USE_SPACES:
+    import spaces
+
 import gradio as gr
 
 import torch
@@ -8,8 +14,11 @@ import time
 
 import gradio as gr
 
-from backbone import extract_features
-from ncut_pytorch import NCUT, rgb_from_tsne_3d, rgb_from_umap_3d
+if USE_SPACES:
+    from backbone import extract_features
+else:
+    from draft_gradio_backbone import extract_features
+from ncut_pytorch import NCUT, eigenvector_to_rgb
 
 
 def compute_ncut(
@@ -24,8 +33,17 @@ def compute_ncut(
     perplexity=150,
     n_neighbors=150,
     min_dist=0.1,
+    sampling_method="fps",
 ):        
     logging_str = ""
+    
+    num_nodes = np.prod(features.shape[:3])
+    if num_nodes / 2 < num_eig:
+        # raise gr.Error("Number of eigenvectors should be less than half the number of nodes.")
+        gr.Warning("Number of eigenvectors should be less than half the number of nodes.\n" f"Setting num_eig to {num_nodes // 2 - 1}.")
+        num_eig = num_nodes // 2 - 1
+        logging_str += f"Number of eigenvectors should be less than half the number of nodes.\n" f"Setting num_eig to {num_nodes // 2 - 1}.\n"
+    
     start = time.time()
     eigvecs, eigvals = NCUT(
         num_eig=num_eig,
@@ -33,33 +51,24 @@ def compute_ncut(
         device="cuda" if torch.cuda.is_available() else "cpu",
         affinity_focal_gamma=affinity_focal_gamma,
         knn=knn_ncut,
+        sample_method=sampling_method,
     ).fit_transform(features.reshape(-1, features.shape[-1]))
     # print(f"NCUT time: {time.time() - start:.2f}s")
     logging_str += f"NCUT time: {time.time() - start:.2f}s\n"
     
     start = time.time()
-    if embedding_method == "UMAP":
-        X_3d, rgb = rgb_from_umap_3d(
-            eigvecs,
-            n_neighbors=n_neighbors,
-            min_dist=min_dist,
-            device="cuda" if torch.cuda.is_available() else "cpu",
-        )
-        # print(f"UMAP time: {time.time() - start:.2f}s")
-        logging_str += f"UMAP time: {time.time() - start:.2f}s\n"
-    elif embedding_method == "t-SNE":    
-        X_3d, rgb = rgb_from_tsne_3d(
-            eigvecs,
-            num_sample=num_sample_tsne,
-            perplexity=perplexity,
-            knn=knn_tsne,
-            device="cuda" if torch.cuda.is_available() else "cpu",
-        )
-        # print(f"t-SNE time: {time.time() - start:.2f}s")
-        logging_str += f"t-SNE time: {time.time() - start:.2f}s\n"
-    else:
-        raise ValueError(f"Embedding method {embedding_method} not supported.")
-    
+    _, rgb = eigenvector_to_rgb(
+        eigvecs,
+        method=embedding_method,
+        num_sample=num_sample_tsne,
+        perplexity=perplexity,
+        n_neighbors=n_neighbors,
+        min_distance=min_dist,
+        knn=knn_tsne,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+    )
+    logging_str += f"{embedding_method} time: {time.time() - start:.2f}s\n"
+
     rgb = rgb.reshape(features.shape[:3] + (3,))
     return rgb, logging_str
 
@@ -76,7 +85,7 @@ def dont_use_too_much_green(image_rgb):
 
 def to_pil_images(images):
     return [
-        Image.fromarray((image * 255).cpu().numpy().astype(np.uint8)).resize((256, 256), Image.NEAREST)
+        Image.fromarray((image * 255).cpu().numpy().astype(np.uint8)).resize((256, 256), Image.Resampling.NEAREST)
         for image in images
     ]
 
@@ -103,11 +112,13 @@ def ncut_run(
     perplexity=500,
     n_neighbors=500,
     min_dist=0.1,
+    sampling_method="fps",
 ):
     logging_str = ""
     if perplexity >= num_sample_tsne or n_neighbors >= num_sample_tsne:
         # raise gr.Error("Perplexity must be less than the number of samples for t-SNE.")
-        gr.Warning("Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting to {num_sample_tsne-1}.")
+        gr.Warning("Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {num_sample_tsne-1}.")
+        logging_str += f"Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {num_sample_tsne-1}.\n"
         perplexity = num_sample_tsne - 1
         n_neighbors = num_sample_tsne - 1
     
@@ -135,26 +146,48 @@ def ncut_run(
         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)
     return to_pil_images(rgb), logging_str
 
-@spaces.GPU(duration=13)
-def quick_run(images, **kwargs):
-    return ncut_run(images, **kwargs)
+def _ncut_run(*args, **kwargs):
+    try:
+        return ncut_run(*args, **kwargs)
+    except Exception as e:
+        gr.Error(str(e))
+        return [], "Error: " + str(e)
 
-@spaces.GPU(duration=30)
-def long_run(images, **kwargs):
-    return ncut_run(images, **kwargs)
+if USE_SPACES:
+    @spaces.GPU(duration=13)
+    def quick_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
 
-@spaces.GPU(duration=60)
-def longer_run(images, **kwargs):
-    return ncut_run(images, **kwargs)
+    @spaces.GPU(duration=30)
+    def long_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
 
-@spaces.GPU(duration=120)
-def super_duper_long_run(images, **kwargs):
-    return ncut_run(images, **kwargs)
+    @spaces.GPU(duration=60)
+    def longer_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
+
+    @spaces.GPU(duration=120)
+    def super_duper_long_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
+
+if not USE_SPACES:
+    def quick_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
+
+    def long_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
+
+    def longer_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
+
+    def super_duper_long_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
 
 def run_fn(
     images,
@@ -171,11 +204,15 @@ def run_fn(
     perplexity=500,
     n_neighbors=500,
     min_dist=0.1,
+    sampling_method="fps",
 ):
     if images is None:
         gr.Warning("No images selected.")
         return [], "No images selected."
     
+    if sampling_method == "fps":
+        sampling_method = "farthest"
+    
     kwargs = {
         "model_name": model_name,
         "layer": layer,
@@ -190,6 +227,7 @@ def run_fn(
         "perplexity": perplexity,
         "n_neighbors": n_neighbors,
         "min_dist": min_dist,
+        "sampling_method": sampling_method,
     }
     num_images = len(images)
     if num_images > 100:
@@ -216,34 +254,48 @@ with gr.Blocks() as demo:
             submit_button = gr.Button("🔴RUN", elem_id="submit_button")
             clear_images_button = gr.Button("🗑️Clear", elem_id='clear_button')
         
-            gr.Markdown('### Load Examples 👇')
-            load_images_button = gr.Button("Load", elem_id="load-images-button")
-            hide_button = gr.Button("Hide", elem_id="hide-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',
+                    '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)"], label="Model", value="SAM(sam_vit_b)", elem_id="model_name")
-            layer_slider = gr.Slider(0, 11, step=1, label="Layer", value=11, elem_id="layer")
-            num_eig_slider = gr.Slider(1, 1000, step=1, label="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="Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for shaper NCUT")
+            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("Additional Parameters", open=False):
-                node_type_dropdown = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Node type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?")
-                num_sample_ncut_slider = gr.Slider(100, 50000, step=100, label="num_sample (NCUT)", value=10000, elem_id="num_sample_ncut", info="Nyström approximation")
-                knn_ncut_slider = gr.Slider(1, 100, step=1, label="KNN (NCUT)", value=10, elem_id="knn_ncut", info="Nyström approximation")
-                embedding_method_dropdown = gr.Dropdown(["t-SNE", "UMAP"], label="Embedding method", value="t-SNE", elem_id="embedding_method")
-                num_sample_tsne_slider = gr.Slider(100, 1000, step=100, label="num_sample (t-SNE/UMAP)", value=300, elem_id="num_sample_tsne", info="Nyström approximation")
-                knn_tsne_slider = gr.Slider(1, 100, step=1, label="KNN (t-SNE/UMAP)", value=10, elem_id="knn_tsne", info="Nyström approximation")
-                perplexity_slider = gr.Slider(10, 500, step=10, label="Perplexity (t-SNE)", value=150, elem_id="perplexity")
-                n_neighbors_slider = gr.Slider(10, 500, step=10, label="n_neighbors (UMAP)", value=150, elem_id="n_neighbors")
-                min_dist_slider = gr.Slider(0.1, 1, step=0.1, label="min_dist (UMAP)", value=0.1, elem_id="min_dist")
+            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")
@@ -253,19 +305,40 @@ with gr.Blocks() as demo:
 
     def empty_input_and_output():
         return [], []
+
+    def load_dataset_images(dataset_name, num_images=10, random_seed=42):
+        from datasets import load_dataset
+        try:
+            dataset = load_dataset(dataset_name)['train']
+        except Exception as e:
+            gr.Error(f"Error loading dataset {dataset_name}: {e}")
+            return None
+        if num_images > len(dataset):
+            num_images = len(dataset)
+        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])
+    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, 
             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
+            perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
         ],
         outputs=[output_gallery, logging_text]
     )
 
 
-demo.launch()
+if USE_SPACES:
+    demo.launch()
+else:
+    demo.launch(share=True)
+
+# %%

requirements.txt

@@ -2,7 +2,7 @@ torch
 torchvision 
 ncut-pytorch
 transformers
+datasets
 segment-anything @ git+https://github.com/facebookresearch/segment-anything.git
 mobile-sam @ git+https://github.com/ChaoningZhang/MobileSAM.git
-
 timm