update color and tsne metric

app.py

@@ -30,7 +30,8 @@ import threading
 
 from ncut_pytorch.backbone import extract_features, load_model
 from ncut_pytorch.backbone import MODEL_DICT, LAYER_DICT, RES_DICT
-from ncut_pytorch import NCUT, eigenvector_to_rgb
+from ncut_pytorch import NCUT
+from ncut_pytorch import eigenvector_to_rgb, rotate_rgb_cube
 
 DATASET_TUPS = [
     # (name, num_classes)
@@ -66,6 +67,7 @@ def compute_ncut(
     knn_ncut=10,
     knn_tsne=10,
     embedding_method="UMAP",
+    embedding_metric='euclidean',
     num_sample_tsne=300,
     perplexity=150,
     n_neighbors=150,
@@ -100,6 +102,7 @@ def compute_ncut(
     _, rgb = eigenvector_to_rgb(
         eigvecs,
         method=embedding_method,
+        metric=embedding_metric,
         num_sample=num_sample_tsne,
         perplexity=perplexity,
         n_neighbors=n_neighbors,
@@ -123,14 +126,20 @@ def dont_use_too_much_green(image_rgb):
     return image_rgb
 
 
-def to_pil_images(images, target_size=256):
+def to_pil_images(images, target_size=512, resize=True):
     size = images[0].shape[1]
     multiplier = target_size // size
     res = int(size * multiplier)
-    return [
-        Image.fromarray((image * 255).cpu().numpy().astype(np.uint8)).resize((res, res), Image.Resampling.NEAREST)
-        for image in images
-    ]
+    pil_images = [
+            Image.fromarray((image * 255).cpu().numpy().astype(np.uint8))
+            for image in images
+        ]
+    if resize:
+        pil_images = [
+            image.resize((res, res), Image.Resampling.NEAREST)
+            for image in pil_images
+        ]
+    return pil_images
     
 
 
@@ -221,6 +230,7 @@ def ncut_run(
     num_sample_ncut=10000,
     knn_ncut=10,
     embedding_method="UMAP",
+    embedding_metric='euclidean',
     num_sample_tsne=1000,
     knn_tsne=10,
     perplexity=500,
@@ -304,6 +314,7 @@ def ncut_run(
                 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,
@@ -333,6 +344,7 @@ def ncut_run(
                 knn_tsne=10,
                 num_sample_tsne=300,
                 embedding_method=embedding_method,
+                embedding_metric=embedding_metric,
                 perplexity=perplexity,
                 n_neighbors=n_neighbors,
                 min_dist=min_dist,
@@ -351,6 +363,7 @@ def ncut_run(
             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,
@@ -378,7 +391,7 @@ def ncut_run(
             galleries = []
             for i_prompt in range(len(lisa_prompts)):
                 _rgb = rgb[i_prompt]
-                galleries.append(to_pil_images(_rgb, target_size=256))
+                galleries.append(to_pil_images(_rgb))
             return *galleries, logging_str
         
         rgb = dont_use_too_much_green(rgb)
@@ -584,6 +597,7 @@ def run_fn(
     num_sample_ncut=10000,
     knn_ncut=10,
     embedding_method="UMAP",
+    embedding_metric='euclidean',
     num_sample_tsne=1000,
     knn_tsne=10,
     perplexity=500,
@@ -699,6 +713,7 @@ def run_fn(
         "num_sample_ncut": num_sample_ncut,
         "knn_ncut": knn_ncut,
         "embedding_method": embedding_method,
+        "embedding_metric": embedding_metric,
         "num_sample_tsne": num_sample_tsne,
         "knn_tsne": knn_tsne,
         "perplexity": perplexity,
@@ -882,9 +897,61 @@ def make_dataset_images_section(advanced=False, is_random=False):
     
     return dataset_dropdown, num_images_slider, random_seed_slider, load_images_button
     
+
+def random_rotate_rgb_gallery(images):
+    if images is None or len(images) == 0:
+        gr.Warning("No images selected.")
+        return []
+    # read webp images
+    images = [Image.open(image[0]).convert("RGB") for image in images]
+    images = [np.array(image).astype(np.float32) for image in images]
+    images = np.stack(images)
+    images = torch.tensor(images) / 255
+    position = np.random.choice([1, 2, 4, 5, 6])
+    images = rotate_rgb_cube(images, position)
+    images = to_pil_images(images, resize=False)
+    return images
+
+def sequence_rotate_rgb_gallery(images):
+    if images is None or len(images) == 0:
+        gr.Warning("No images selected.")
+        return []
+    # read webp images
+    images = [Image.open(image[0]).convert("RGB") for image in images]
+    images = [np.array(image).astype(np.float32) for image in images]
+    images = np.stack(images)
+    images = torch.tensor(images) / 255
+    rotation_matrix = torch.tensor([[0, 1, 0], [0, 0, 1], [1, 0, 0]]).float()
+    images = images @ rotation_matrix
+    images = to_pil_images(images, resize=False)
+    return images
+
+def flip_rgb_gallery(images, axis=0):
+    if images is None or len(images) == 0:
+        gr.Warning("No images selected.")
+        return []
+    # read webp images
+    images = [Image.open(image[0]).convert("RGB") for image in images]
+    images = [np.array(image).astype(np.float32) for image in images]
+    images = np.stack(images)
+    images = torch.tensor(images) / 255
+    images = 1 - images
+    images = to_pil_images(images, resize=False)
+    return images
+
+def add_output_images_buttons(output_gallery):
+    with gr.Row():
+        rotate_button = gr.Button("🔄 Rotate", elem_id="rotate_button", variant='secondary')
+        rotate_button.click(sequence_rotate_rgb_gallery, inputs=[output_gallery], outputs=[output_gallery])
+        flip_button = gr.Button("🔃 Flip", elem_id="flip_button", variant='secondary')
+        flip_button.click(flip_rgb_gallery, inputs=[output_gallery], outputs=[output_gallery])
+    return rotate_button, flip_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", show_share_button=True)
+    output_gallery = gr.Gallery(format='png', value=[], label="NCUT Embedding", show_label=False, elem_id="ncut", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True)
+    add_output_images_buttons(output_gallery)
     return output_gallery
 
 def make_parameters_section(is_lisa=False):
@@ -951,6 +1018,7 @@ def make_parameters_section(is_lisa=False):
         sampling_method_dropdown = gr.Dropdown(["fps", "random"], label="NCUT: Sampling method", value="fps", elem_id="sampling_method", info="Nyström approximation")
         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")
+        embedding_metric_dropdown = gr.Dropdown(["euclidean", "cosine"], label="t-SNE/UMAP metric", value="euclidean", elem_id="embedding_metric")
         num_sample_tsne_slider = gr.Slider(100, 10000, 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, 1000, step=10, label="t-SNE: perplexity", value=150, elem_id="perplexity")
@@ -958,7 +1026,7 @@ def make_parameters_section(is_lisa=False):
         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, 
+            embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
             perplexity_slider, n_neighbors_slider, min_dist_slider, 
             sampling_method_dropdown, positive_prompt, negative_prompt]
 
@@ -980,7 +1048,7 @@ with demo:
                 [
                     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, 
+                    embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
                     sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
@@ -999,7 +1067,7 @@ with demo:
                 positive_prompt, negative_prompt,
                 false_placeholder, no_prompt, no_prompt, no_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
-                embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
             ],
             outputs=[output_gallery, logging_text],
@@ -1039,7 +1107,7 @@ with demo:
                 [
                     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, 
+                    embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
                     sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
@@ -1062,7 +1130,7 @@ with demo:
                 positive_prompt, negative_prompt,
                 false_placeholder, no_prompt, no_prompt, no_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
-                embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown, 
                 old_school_ncut_checkbox
             ],
@@ -1079,13 +1147,16 @@ with demo:
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #1)')
-                l1_gallery = gr.Gallery(value=[], label="Recursion #1", show_label=False, elem_id="ncut_l1", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
+                l1_gallery = gr.Gallery(format='png', value=[], label="Recursion #1", show_label=False, elem_id="ncut_l1", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
+                add_output_images_buttons(l1_gallery)
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #2)')
-                l2_gallery = gr.Gallery(value=[], label="Recursion #2", show_label=False, elem_id="ncut_l2", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
+                l2_gallery = gr.Gallery(format='png', value=[], label="Recursion #2", show_label=False, elem_id="ncut_l2", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
+                add_output_images_buttons(l2_gallery)
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #3)')
-                l3_gallery = gr.Gallery(value=[], label="Recursion #3", show_label=False, elem_id="ncut_l3", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
+                l3_gallery = gr.Gallery(format='png', value=[], label="Recursion #3", show_label=False, elem_id="ncut_l3", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
+                add_output_images_buttons(l3_gallery)
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 input_gallery, submit_button, clear_images_button = make_input_images_section()
@@ -1105,7 +1176,7 @@ with demo:
                 [
                     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, 
+                    embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
                     sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
@@ -1136,7 +1207,7 @@ with demo:
                 positive_prompt, negative_prompt,
                 false_placeholder, no_prompt, no_prompt, no_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
-                embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown,
                 false_placeholder, number_placeholder, true_placeholder,
                 l2_num_eig_slider, l3_num_eig_slider, metric_dropdown, 
@@ -1156,7 +1227,7 @@ with demo:
                 [
                     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, 
+                    embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
                     sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
@@ -1179,7 +1250,7 @@ with demo:
                 positive_prompt, negative_prompt,
                 false_placeholder, no_prompt, no_prompt, no_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
-                embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown,
                 place_holder_false, max_frame_number
             ],
@@ -1204,15 +1275,15 @@ with demo:
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Prompt #1)')
-                l1_gallery = gr.Gallery(value=[], label="Prompt #1", show_label=False, elem_id="ncut_p1", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
+                l1_gallery = gr.Gallery(format='png', value=[], label="Prompt #1", show_label=False, elem_id="ncut_p1", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
                 prompt1 = gr.Textbox(label="Input Prompt #1", elem_id="prompt1", value="where is the person, include the clothes, don't include the guitar and chair", lines=3)
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Prompt #2)')
-                l2_gallery = gr.Gallery(value=[], label="Prompt #2", show_label=False, elem_id="ncut_p2", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
+                l2_gallery = gr.Gallery(format='png', value=[], label="Prompt #2", show_label=False, elem_id="ncut_p2", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
                 prompt2 = gr.Textbox(label="Input Prompt #2", elem_id="prompt2", value="where is the Gibson Les Pual guitar", lines=3)
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Prompt #3)')
-                l3_gallery = gr.Gallery(value=[], label="Prompt #3", show_label=False, elem_id="ncut_p3", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
+                l3_gallery = gr.Gallery(format='png', value=[], label="Prompt #3", show_label=False, elem_id="ncut_p3", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True)
                 prompt3 = gr.Textbox(label="Input Prompt #3", elem_id="prompt3", value="where is the floor", lines=3)
                 
         with gr.Row():
@@ -1225,7 +1296,7 @@ with demo:
                 [
                     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, 
+                    embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
                     sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section(is_lisa=True)
@@ -1240,7 +1311,7 @@ with demo:
                 positive_prompt, negative_prompt,
                 true_placeholder, prompt1, prompt2, prompt3,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
-                embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
             ],
             outputs=galleries + [logging_text],
@@ -1273,7 +1344,7 @@ with demo:
                 [
                     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, 
+                    embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
                     sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
@@ -1309,7 +1380,7 @@ with demo:
                 positive_prompt, negative_prompt,
                 false_placeholder, no_prompt, no_prompt, no_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
-                embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
             ],
             # outputs=galleries + [logging_text],
@@ -1320,12 +1391,13 @@ with demo:
         def add_one_model(i_model=1):
             with gr.Column(scale=5, min_width=200) as col:
                 gr.Markdown(f'### Output Images')
-                output_gallery = gr.Gallery(value=[], label="NCUT Embedding", show_label=False, elem_id=f"ncut{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True)
+                output_gallery = gr.Gallery(format='png', value=[], label="NCUT Embedding", show_label=False, elem_id=f"ncut{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True)
                 submit_button = gr.Button("🔴 RUN", elem_id=f"submit_button{i_model}", variant='primary')
+                add_output_images_buttons(output_gallery)
                 [
                     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, 
+                    embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
                     sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
@@ -1341,7 +1413,7 @@ with demo:
                         positive_prompt, negative_prompt,
                         false_placeholder, no_prompt, no_prompt, no_prompt,
                         affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
-                        embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                        embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                         perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
                     ],
                     outputs=[output_gallery, logging_text]
@@ -1390,14 +1462,18 @@ with demo:
     
     
     with gr.Tab('📄About'):
-        gr.Markdown("##### This demo is for python package `ncut-pytorch`, please visit the [Documentation](https://ncut-pytorch.readthedocs.io/) ")
+        gr.Markdown("**This demo is for the Python package `ncut-pytorch`, please visit the [Documentation](https://ncut-pytorch.readthedocs.io/)**")
+        gr.Markdown("**All the models and functions used for this demo are in the Python package `ncut-pytorch`**")
         gr.Markdown("---")
         gr.Markdown("**Normalized Cuts**, aka. spectral clustering, is a graphical method to analyze data grouping in the affinity eigenvector space. It has been widely used for unsupervised segmentation in the 2000s.")
+        gr.Markdown("*Normalized Cuts and Image Segmentation, Jianbo Shi and Jitendra Malik, 2000*")
         gr.Markdown("---")        
-        gr.Markdown("##### We have implemented NCut, with some advanced features:")
+        gr.Markdown("**We have improved NCut, with some advanced features:**")
         gr.Markdown("- **Nyström** Normalized Cut, is a new approximation algorithm developed for large-scale graph cuts, a large-graph of million nodes can be processed in under 10s (cpu) or 2s (gpu).")
         gr.Markdown("- **spectral-tSNE** visualization, a new method to visualize the high-dimensional eigenvector space with 3D RGB cube. Color is aligned across images, color infers distance in representation.")
-        
+        gr.Markdown("*paper in prep, Yang 2024*")
+        gr.Markdown("*AlignedCut: Visual Concepts Discovery on Brain-Guided Universal Feature Space, Huzheng Yang, James Gee\*, and Jianbo Shi\*, 2024*")
+
         
     with gr.Row():
         with gr.Column():