add progress bar

app.py

@@ -77,7 +77,9 @@ def compute_ncut(
     min_dist=0.1,
     sampling_method="fps",
     metric="cosine",
+    progess_start=0.4,
 ):        
+    progress = gr.Progress()
     logging_str = ""
     
     num_nodes = np.prod(features.shape[:-1])
@@ -88,6 +90,7 @@ def compute_ncut(
         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()
+    progress(progess_start+0.0, desc="NCut")
     eigvecs, eigvals = NCUT(
         num_eig=num_eig,
         num_sample=num_sample_ncut,
@@ -102,6 +105,7 @@ def compute_ncut(
     logging_str += f"NCUT time: {time.time() - start:.2f}s\n"
     
     start = time.time()
+    progress(progess_start+0.01, desc="spectral-tSNE")
     _, rgb = eigenvector_to_rgb(
         eigvecs,
         method=embedding_method,
@@ -249,15 +253,34 @@ def blend_image_with_heatmap(image, heatmap, opacity1=0.5, opacity2=0.5):
     blended = (1 - opacity1) * image + opacity2 * heatmap
     return blended.astype(np.uint8)
 
-def make_cluster_plot(eigvecs, images, h=64, w=64):
+def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6):
+    progress = gr.Progress()
+    progress(progess_start, desc="Finding Clusters by FPS")
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    eigvecs = eigvecs.to(device)
     from ncut_pytorch.ncut_pytorch import farthest_point_sampling
     magnitude = torch.norm(eigvecs, dim=-1)
-    p = 0.5
+    p = 0.8
     top_p_idx = magnitude.argsort(descending=True)[:int(p * magnitude.shape[0])]
-    num_samples = 50
+    num_samples = 300
+    if num_samples > top_p_idx.shape[0]:
+        num_samples = top_p_idx.shape[0]
     fps_idx = farthest_point_sampling(eigvecs[top_p_idx], num_samples)
     fps_idx = top_p_idx[fps_idx]
     
+    # fps round 2 on the heatmap
+    left = eigvecs[fps_idx, :].clone()
+    right = eigvecs.clone()
+    left = F.normalize(left, dim=-1)
+    right = F.normalize(right, dim=-1)
+    heatmap = left @ right.T
+    heatmap = F.normalize(heatmap, dim=-1)
+    num_samples = 80
+    if num_samples > fps_idx.shape[0]:
+        num_samples = fps_idx.shape[0]
+    r2_fps_idx = farthest_point_sampling(heatmap, num_samples)
+    fps_idx = fps_idx[r2_fps_idx]
+    
     # downsample to 256x256
     images = F.interpolate(images, (256, 256), mode="bilinear")
     images = images.cpu().numpy()
@@ -269,29 +292,57 @@ def make_cluster_plot(eigvecs, images, h=64, w=64):
     # sort the fps_idx by the mean of the heatmap
     fps_heatmaps = {}
     sort_values = []
+    top3_image_idx = {}
     for _, idx in enumerate(fps_idx):
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        eigvecs = eigvecs.to(device)
         heatmap = F.cosine_similarity(eigvecs, eigvecs[idx][None], dim=-1)
+        
+        # def top_percentile(tensor, p=0.8, max_size=10000):
+        #     tensor = tensor.clone().flatten()
+        #     if tensor.shape[0] > max_size:
+        #         tensor = tensor[torch.randperm(tensor.shape[0])[:max_size]]
+        #     return tensor.quantile(p)
+        # top_p = top_percentile(heatmap, p=0.5)
+        top_p = 0.5
+        
         heatmap = heatmap.reshape(-1, h, w)
-        mask = (heatmap > 0.5).float()    
+        mask = (heatmap > top_p).float()
+        # take top 3 masks only
+        mask_sort_values = mask.mean((1, 2))
+        mask_sort_idx = torch.argsort(mask_sort_values, descending=True)
+        mask = mask[mask_sort_idx[:3]] 
         sort_values.append(mask.mean().item())
-        fps_heatmaps[idx.item()] = heatmap.cpu()
+        # fps_heatmaps[idx.item()] = heatmap.cpu()
+        fps_heatmaps[idx.item()] = heatmap[mask_sort_idx[:3]].cpu()
+        top3_image_idx[idx.item()] = mask_sort_idx[:3]
+    # do the sorting
+    _sort_idx = torch.tensor(sort_values).argsort(descending=True)
+    fps_idx = fps_idx[_sort_idx]
+    # reverse the fps_idx
+    # fps_idx = fps_idx.flip(0)
+    # discard the big clusters
+    fps_idx = fps_idx[10:]
+    # shuffle the fps_idx
+    fps_idx = fps_idx[torch.randperm(fps_idx.shape[0])]
     
     fig_images = []
     i_cluster = 0
-    for i_fig in range(10):    
+    num_plots = 10
+    plot_step_float = (1.0 - progess_start) / num_plots
+    for i_fig in range(num_plots):    
+        progress(progess_start + i_fig * plot_step_float, desc="Plotting Clusters")
         fig, axs = plt.subplots(3, 5, figsize=(15, 9))
         for ax in axs.flatten():
             ax.axis("off")
         for j, idx in enumerate(fps_idx[i_fig*5:i_fig*5+5]):
             heatmap = fps_heatmaps[idx.item()]
-            mask = (heatmap > 0.1).float()
-            sorted_image_idxs = torch.argsort(mask.mean((1, 2)), descending=True)
+            # mask = (heatmap > 0.1).float()
+            # sorted_image_idxs = torch.argsort(mask.mean((1, 2)), descending=True)
             size = (images.shape[1], images.shape[2])
             heatmap = apply_reds_colormap(heatmap, size)
-            for i, image_idx in enumerate(sorted_image_idxs[:3]):
-                _heatmap = blend_image_with_heatmap(images[image_idx], heatmap[image_idx])
+            # for i, image_idx in enumerate(sorted_image_idxs[:3]):
+            for i, image_idx in enumerate(top3_image_idx[idx.item()]):
+                # _heatmap = blend_image_with_heatmap(images[image_idx], heatmap[image_idx])
+                _heatmap = blend_image_with_heatmap(images[image_idx], heatmap[i])
                 axs[i, j].imshow(_heatmap)
                 if i == 0:
                     axs[i, j].set_title(f"cluster {i_cluster+1}", fontsize=24)
@@ -348,6 +399,9 @@ def ncut_run(
     lisa_prompt2="",
     lisa_prompt3="",
 ):
+    progress = gr.Progress()
+    progress(0.2, desc="Feature Extraction")
+    
     logging_str = ""
     if "AlignedThreeModelAttnNodes" == model_name:
         # dirty patch for the alignedcut paper
@@ -396,12 +450,16 @@ 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"
     
+    progress(0.4, desc="NCut")
+    
     if recursion:
         rgbs = []
         recursion_gammas = [recursion_l1_gamma, recursion_l2_gamma, recursion_l3_gamma]
         inp = features
+        progress_start = 0.4
         for i, n_eigs in enumerate([num_eig, recursion_l2_n_eigs, recursion_l3_n_eigs]):
             logging_str += f"Recursion #{i+1}\n"
+            progress_start += + 0.1 * i
             rgb, _logging_str, eigvecs = compute_ncut(
                 inp,
                 num_eig=n_eigs,
@@ -417,6 +475,7 @@ def ncut_run(
                 min_dist=min_dist,
                 sampling_method=sampling_method,
                 metric="cosine" if i == 0 else recursion_metric,
+                progess_start=progress_start,
             )
             logging_str += _logging_str
             
@@ -424,6 +483,7 @@ def ncut_run(
             if "AlignedThreeModelAttnNodes" == model_name:
                 # dirty patch for the alignedcut paper
                 start = time.time()
+                progress(progress_start + 0.09, desc=f"Plotting Recursion {i+1}")
                 pil_images = []
                 for i_image in range(rgb.shape[0]):
                     _im = plot_one_image_36_grid(images[i_image], rgb[i_image])
@@ -442,6 +502,8 @@ def ncut_run(
     if old_school_ncut:  # individual images
         logging_str += "Running NCut for each image independently\n"
         rgb = []
+        progress_start = 0.4
+        step_float = 0.6 / features.shape[0]
         for i_image in range(features.shape[0]):
             logging_str += f"Image #{i_image+1}\n"
             feature = features[i_image]
@@ -459,6 +521,7 @@ def ncut_run(
                 n_neighbors=n_neighbors,
                 min_dist=min_dist,
                 sampling_method=sampling_method,
+                progess_start=progress_start+step_float*i_image,
             )
             logging_str += _logging_str
             rgb.append(_rgb[0])
@@ -486,6 +549,7 @@ def ncut_run(
         if "AlignedThreeModelAttnNodes" == model_name:
             # dirty patch for the alignedcut paper
             start = time.time()
+            progress(0.6, desc="Plotting")
             pil_images = []
             for i_image in range(rgb.shape[0]):
                 _im = plot_one_image_36_grid(images[i_image], rgb[i_image])
@@ -506,15 +570,18 @@ def ncut_run(
         
         if not video_output:
             start = time.time()
+            progress_start = 0.6
+            progress(progress_start, desc="Plotting Clusters")
             h, w = features.shape[1], features.shape[2]
             if torch.cuda.is_available():
                 images = images.cuda()
             _images = reverse_transform_image(images, stablediffusion="stable" in model_name.lower())
-            cluster_images = make_cluster_plot(eigvecs, _images, h=h, w=w)
+            cluster_images = make_cluster_plot(eigvecs, _images, h=h, w=w, progess_start=progress_start)
             logging_str += f"plot time: {time.time() - start:.2f}s\n"
         
     
     if video_output:
+        progress(0.8, desc="Saving Video")
         video_path = get_random_path()
         video_cache.add_video(video_path)
         pil_images_to_video(to_pil_images(rgb), video_path)
@@ -526,26 +593,26 @@ def ncut_run(
 
 def _ncut_run(*args, **kwargs):
     n_ret = kwargs.pop("n_ret", 1)
-    # try:
-    #     if torch.cuda.is_available():
-    #         torch.cuda.empty_cache()
+    try:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
             
-    #     ret = ncut_run(*args, **kwargs)
+        ret = ncut_run(*args, **kwargs)
         
-    #     if torch.cuda.is_available():
-    #         torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         
-    #     ret = list(ret)[:n_ret] + [ret[-1]]
-    #     return ret
-    # except Exception as e:
-    #     gr.Error(str(e))
-    #     if torch.cuda.is_available():
-    #         torch.cuda.empty_cache()
-    #     return *(None for _ in range(n_ret)), "Error: " + str(e)
-
-    ret = ncut_run(*args, **kwargs)
-    ret = list(ret)[:n_ret] + [ret[-1]]
-    return ret
+        ret = list(ret)[:n_ret] + [ret[-1]]
+        return ret
+    except Exception as e:
+        gr.Error(str(e))
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return *(None for _ in range(n_ret)), "Error: " + str(e)
+
+    # ret = ncut_run(*args, **kwargs)
+    # ret = list(ret)[:n_ret] + [ret[-1]]
+    # return ret
 
 if USE_HUGGINGFACE_ZEROGPU:
     @spaces.GPU(duration=20)
@@ -744,10 +811,15 @@ def run_fn(
     n_ret=1,
 ):
     
+    progress=gr.Progress()
+    progress(0, desc="Starting")
+    
+    
     if images is None:
         gr.Warning("No images selected.")
         return *(None for _ in range(n_ret)), "No images selected."
     
+    progress(0.05, desc="Processing Images")
     video_output = False
     if isinstance(images, str):
         images = extract_video_frames(images, max_frames=max_frames)
@@ -767,6 +839,7 @@ def run_fn(
     images = [transform_image(image, resolution=resolution, stablediffusion=stablediffusion) for image in images]
     images = torch.stack(images)
     
+    progress(0.1, desc="Downloading Model")
     
     if is_lisa:
         import subprocess
@@ -976,10 +1049,13 @@ def make_dataset_images_section(advanced=False, is_random=False):
     def load_dataset_images(is_advanced, dataset_name, num_images=10, 
                             is_filter=True, filter_by_class_text="0,1,2", 
                             is_random=False, seed=1):
+        progress = gr.Progress()
+        progress(0, desc="Loading Images")
         if is_advanced == "Basic":
             gr.Info("Loaded images from Ego-Exo4D")
             return default_images
         try:
+            progress(0.5, desc="Downloading Dataset")
             dataset = load_dataset(dataset_name, trust_remote_code=True)
             key = list(dataset.keys())[0]
             dataset = dataset[key]
@@ -990,6 +1066,7 @@ def make_dataset_images_section(advanced=False, is_random=False):
             num_images = len(dataset)
         
         if is_filter:
+            progress(0.8, desc="Filtering Images")
             classes = [int(i) for i in filter_by_class_text.split(",")]
             labels = np.array(dataset['label'])
             unique_labels = np.unique(labels)
@@ -1193,6 +1270,7 @@ with demo:
             with gr.Column(scale=5, min_width=200):
                 input_gallery, submit_button, clear_images_button = make_input_images_section()
                 dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_dataset_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):

app_text.py

@@ -150,6 +150,7 @@ def ncut_run(
     min_dist=0.1,
     sampling_method="fps",
 ):
+    progress = gr.Progress()
     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.")
@@ -163,6 +164,7 @@ def ncut_run(
         
     node_type = node_type.split(":")[0].strip()
     
+    progress(0.5, desc="Feature Extraction")
     model = model.to("cuda" if torch.cuda.is_available() else "cpu")
     
     start = time.time()
@@ -180,6 +182,7 @@ 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"
     
+    progress(0.6, desc="NCUT & spectral-tSNE")
     rgb, _logging_str, _ = compute_ncut(
         features,
         num_eig=num_eig,
@@ -197,6 +200,7 @@ def ncut_run(
     logging_str += _logging_str
     
     start = time.time() 
+    progress(0.8, desc="Plotting")
     title = f"{model_name}, Layer {layer}, {node_type}"
     fig = make_plot(token_texts, rgb, title=title)
     logging_str += f"Plotting time: {time.time() - start:.2f}s\n"
@@ -223,6 +227,8 @@ else:
         return _ncut_run(*args, **kwargs)
 
 def real_run(model_name, text, layer, node_type, num_eig, affinity_focal_gamma, num_sample_ncut, knn_ncut, embedding_method, num_sample_tsne, knn_tsne, perplexity, n_neighbors, min_dist, sampling_method):
+    progress = gr.Progress()
+    progress(0.1, desc="Downloading model")
     model = TEXT_MODEL_DICT[model_name]()
     return __ncut_run(model, text, model_name, layer, num_eig, node_type, 
                 affinity_focal_gamma, num_sample_ncut, knn_ncut, embedding_method, 
@@ -251,7 +257,9 @@ def make_demo():
             clear_button = gr.Button("🗑️Clear", elem_id='clear_button', variant='stop')
         with gr.Column(scale=5, min_width=200):
             gr.Markdown("### Parameters <a style='color: #0044CC;' href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Help</a>")
-            model_name = gr.Dropdown(list(TEXT_MODEL_DICT.keys()), label="Model", value="meta-llama/Meta-Llama-3.1-8B")
+            model_list = list(TEXT_MODEL_DICT.keys())
+            model_list = [model for model in model_list if model != "meta-llama/Meta-Llama-3-8B"]
+            model_name = gr.Dropdown(model_list, label="Model", value="meta-llama/Meta-Llama-3.1-8B")
             layer = gr.Slider(1, 32, step=1, value=32, label="Layer")
             node_type = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Node Type", value="block: sum of residual")
             num_eig = gr.Slider(minimum=1, maximum=1000, step=1, value=100, label="Number of Eigenvectors")