update sub-cluster

app.py

@@ -148,7 +148,7 @@ def compute_ncut(
     logging_str = ""
     
     num_nodes = np.prod(features.shape[:-1])
-    if num_nodes / 2 < num_eig:
+    if num_nodes / 2 - 1 < 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
@@ -2196,7 +2196,7 @@ demo = gr.Blocks(
     css=custom_css,
 )
 with demo:
-
+        
     with gr.Tab('PlayGround'):
         eigvecs = gr.State(np.array([]))
         tsne3d_rgb = gr.State(np.array([]))
@@ -4246,6 +4246,219 @@ with demo:
                             fg_contrast_slider, bg_contrast_slider],
                     outputs=[mask_gallery, crop_gallery])
                 
+                
+    with gr.Tab('Sub-cluster (dev)') as sub_cluster_tab:
+        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 = 300
+                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()
+                add_download_button(output_gallery)
+                with gr.Accordion("Sub-cluster", open=True):
+                    num_image_cluster_slider = gr.Slider(1, 20, step=1, label="Number of clusters (image-level)", value=6, elem_id="num_image_cluster", info="Image-level clustering before pixel-level clustering")
+                    show_class_label_checkbox = gr.Checkbox(label="Show cluster label", value=True, elem_id="show_class_label", info="Show image-level clustering label on the output")
+                    overlay_original_image_checkbox = gr.Checkbox(label="Overlay original image", value=False, elem_id="overlay_original_image", info="Overlay original image on the output")
+                [
+                    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()
+
+        with gr.Row():
+            image_cluster_plot = gr.Image(value=None, label="Image-level clustering", elem_id="image_cluster_plot", interactive=False)
+        
+        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)
+        
+        def cluster_image_level(image_gallery, n_clusters=6):
+            images = [tup[0] for tup in image_gallery]
+            if isinstance(images[0], str):
+                images = [Image.open(image) for image in images]
+            resized_images = [image.resize((224, 224), Image.LANCZOS) for image in images]
+            images = [transform_image(image, resolution=(224, 224), stablediffusion=False) for image in resized_images]
+            images = torch.stack(images)
+            
+            is_cuda = torch.cuda.is_available()
+            
+            dino = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16').eval()
+            from torch.utils.data import DataLoader, TensorDataset
+            dataset = TensorDataset(images)
+            dataloader = DataLoader(dataset, batch_size=128, shuffle=False)
+            features = []
+            if is_cuda:
+                dino = dino.cuda()
+            with torch.no_grad():
+                for _images in dataloader:
+                    _images = _images[0]
+                    _images = _images.cuda() if is_cuda else _images
+                    feature = dino.get_intermediate_layers(_images, 12)[11]
+                    feature = feature[:, 0]
+                    # normalize the feature
+                    feature = torch.nn.functional.normalize(feature, dim=-1)
+                    features.append(feature.cpu())
+            class_features = torch.cat(features, dim=0)
+            
+            from ncut_pytorch import NCUT
+            inp = class_features
+            b, c = inp.shape
+            inp = inp.reshape(b, c)
+            class_ncut = NCUT(num_eig=n_clusters, move_output_to_cpu=True, distance='cosine', knn=1, device='cuda' if is_cuda else 'cpu')
+            class_eigvecs, _ = class_ncut.fit_transform(inp)
+            from ncut_pytorch import kway_ncut
+            kway_eigvecs = kway_ncut(class_eigvecs.cuda())  # (n, k), one-hot
+            kway_indices = torch.argmax(kway_eigvecs, dim=1)
+            cluster_labels = kway_indices.cpu().numpy()
+            
+            from sklearn.manifold import TSNE
+            perplexity = min(30, int(len(cluster_labels)/2 - 1))
+            tsne_2d = TSNE(n_components=2, perplexity=perplexity, random_state=0).fit_transform(class_eigvecs.cpu())
+            
+            tsne_image_plot = plot_2d_tsne_with_images(cluster_labels, tsne_2d, resized_images)
+            
+            return cluster_labels, tsne_image_plot, resized_images
+        
+        def plot_2d_tsne_with_images(cluster_labels, tsne_2d, images, max_num_images=300):
+            # use dark background
+            plt.style.use('dark_background')
+            
+            def pad_image_with_border(image, border_color, border_width):
+                new_image = np.ones((image.shape[0] + 2 * border_width, image.shape[1] + 2 * border_width, image.shape[2]), dtype=image.dtype)
+                new_image[:, :] = border_color
+                new_image[border_width:-border_width, border_width:-border_width] = image
+                return new_image
+
+            padded_images = []
+            for i in range(len(tsne_2d)):
+                image = images[i]
+                image = np.array(image)
+                
+                border_color = np.array(plt.cm.tab20(cluster_labels[i] / 20))[:3] * 255
+                border_width = 20
+                padded_image = pad_image_with_border(image, border_color, border_width)
+                padded_images.append(padded_image)
+                            
+            # Plot the t-SNE points
+            fig, ax = plt.subplots(1, 1, figsize=(15, 15))
+            ax.scatter(tsne_2d[:, 0], tsne_2d[:, 1], s=20, c=cluster_labels, cmap='tab20')
+            ax.set_xticks([])
+            ax.set_yticks([])
+            ax.axis('off')
+            ax.set_xlim(tsne_2d[:, 0].min()*1.1, tsne_2d[:, 0].max()*1.1)
+            ax.set_ylim(tsne_2d[:, 1].min()*1.1, tsne_2d[:, 1].max()*1.1)
+            from matplotlib.offsetbox import AnnotationBbox, OffsetImage
+
+            if len(tsne_2d) > max_num_images:
+                random_indices = np.random.choice(len(tsne_2d), max_num_images, replace=False)
+                tsne_2d = tsne_2d[random_indices]
+                padded_images = [padded_images[i] for i in random_indices]
+                
+
+            # Add the top1_image_blended to the scatter plot
+            for i, (x, y) in enumerate(tsne_2d):
+                img = padded_images[i]
+                img = np.array(img)
+                imgbox = OffsetImage(img, zoom=0.15)
+                ab = AnnotationBbox(imgbox, (x, y), frameon=False)
+                ax.add_artist(ab)
+
+            # Remove the white space around the plot
+            fig.tight_layout(pad=0)
+            
+            # Save the plot to an in-memory buffer
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+            buf.seek(0)
+
+            # Load the image into a NumPy array
+            image = np.array(Image.open(buf))
+
+            # Close the buffer and plot
+            buf.close()
+            plt.close(fig)
+            
+            pil_image = Image.fromarray(image)
+            
+            # Reset the style
+            plt.style.use('default')
+            
+            return pil_image
+        
+        def add_boarder_and_overlay_image(ncut_images, original_images, cluster_labels, is_overlay_image, is_add_border):
+            def pad_image_with_left_border(image, border_color, border_width):
+                new_image = np.ones((image.shape[0], image.shape[1] + border_width, image.shape[2]), dtype=image.dtype)
+                new_image[:, :] = border_color
+                new_image[:, border_width:] = image
+                return new_image
+            
+            if is_overlay_image:
+                for i, ncut_image in enumerate(ncut_images):
+                    original_image = original_images[i]
+                    original_image = np.array(original_image)
+                    ncut_image = ncut_image.resize(original_image.shape[:2][::-1], Image.LANCZOS)
+                    ncut_image = np.array(ncut_image)
+                    blend_image = 0.5 * original_image + 0.5 * ncut_image
+                    blend_image = blend_image.astype(np.uint8)
+                    ncut_images[i] = Image.fromarray(blend_image)
+            if is_add_border:
+                for i, ncut_image in enumerate(ncut_images):
+                    border_color = np.array(plt.cm.tab20(cluster_labels[i] / 20))[:3] * 255
+                    border_width = 20
+                    ncut_image = ncut_image.resize((224, 224), Image.LANCZOS)
+                    ncut_image = np.array(ncut_image)
+                    ncut_image = pad_image_with_left_border(ncut_image, border_color, border_width)
+                    ncut_images[i] = Image.fromarray(ncut_image)
+            
+            return ncut_images
+        
+        def sub_cluster_run_fn(images, n_clusters, is_overlay_image, is_add_border, *args, **kwargs):
+            cluster_labels, tsne_image_plot, resized_images = cluster_image_level(images, n_clusters)
+                        
+            output_images = [None] * len(images)
+            for i in range(n_clusters):
+                indices = np.where(cluster_labels == i)[0]
+                if len(indices) == 0:
+                    continue
+                input_images = [images[j] for j in indices]
+                _i_output_images, logging_text = run_fn(input_images, *args, **kwargs)
+                for _i, _idx in enumerate(indices):
+                    output_images[_idx] = _i_output_images[_i]
+                    
+            output_images = add_boarder_and_overlay_image(output_images, resized_images, cluster_labels, is_overlay_image, is_add_border)
+            return output_images, tsne_image_plot
+        
+        if USE_HUGGINGFACE_ZEROGPU:
+            @spaces.GPU(duration=120)
+            def sub_cluster_run_fn_wrap(*args, **kwargs):
+                return sub_cluster_run_fn(*args, **kwargs)
+        else:
+            sub_cluster_run_fn_wrap = sub_cluster_run_fn
+        
+        # image_cluster_indices = gr.State(np.array([]))    
+        
+        # submit_button.click(
+        #     cluster_image_level,
+        #     inputs=[input_gallery, num_image_cluster_slider],
+        #     outputs=[image_cluster_indices, image_cluster_plot],
+        # )
+        
+        submit_button.click(
+            partial(sub_cluster_run_fn_wrap, n_ret=1, plot_clusters=False),
+            inputs=[
+                input_gallery, num_image_cluster_slider, overlay_original_image_checkbox, show_class_label_checkbox,
+                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
+            ],
+            outputs=[output_gallery, image_cluster_plot],
+        )                
     
     # with gr.Tab('PlayGround (test)', visible=False) as test_playground_tab:
     #     eigvecs = gr.State(np.array([]))
@@ -4569,7 +4782,7 @@ with demo:
         
         hidden_button.change(update_smile, [n_smiles], [n_smiles, hidden_button])
         hidden_tabs = [tab_alignedcut_advanced, tab_model_aligned_advanced, tab_recursivecut_advanced, 
-                        tab_compare_models_advanced, tab_directed_ncut, tab_aligned, tab_lisa]
+                        tab_compare_models_advanced, tab_directed_ncut, tab_aligned, tab_lisa, sub_cluster_tab]
         hidden_button.change(partial(unlock_tabs, n_tab=len(hidden_tabs)), [n_smiles], hidden_tabs)
                 
     with gr.Row():