add tree+image

app.py

@@ -16,6 +16,7 @@ import multiprocessing as mp
 from einops import rearrange
 from matplotlib import pyplot as plt
 import matplotlib
+from matplotlib.offsetbox import AnnotationBbox, OffsetImage
 USE_HUGGINGFACE_ZEROGPU = os.getenv("USE_HUGGINGFACE_ZEROGPU", "False").lower() in ["true", "1", "yes"]
 DOWNLOAD_ALL_MODELS_DATASETS = os.getenv("DOWNLOAD_ALL_MODELS_DATASETS", "False").lower() in ["true", "1", "yes"]
 
@@ -2342,13 +2343,99 @@ with demo:
                     
                     pil_image = Image.fromarray(image)
                     return pil_image
-                                
-                def run_fps_tsne_hierarchical(eigvecs, num_sample_fps, perplexity_tsne, tsne3d_rgb, seed=0):
+                
+                def get_top1_heatmap_for_each_dot(images, eigvecs, fps_eigvecs, max_display_dots, fps_tsne_rgb, tsne_embed):
+                    n_dots = fps_eigvecs.shape[0]
+                    if n_dots > max_display_dots:
+                        dots_idx = np.random.choice(n_dots, max_display_dots, replace=False)
+                        import fpsample
+                        dots_idx = fpsample.bucket_fps_kdline_sampling(tsne_embed, max_display_dots, 5).astype(np.int64)
+                    else:
+                        dots_idx = np.arange(n_dots)
+                    fps_eigvecs = fps_eigvecs[dots_idx]
+                    fps_tsne_rgb = fps_tsne_rgb[dots_idx]
+                    
+                    heatmaps = eigvecs @ fps_eigvecs.T  # [B, H, W, C] @ [N, C] -> [B, H, W, N]
+                    value = heatmaps.mean(1).mean(1)  # [B, N]
+                    top1_image_idxs = value.argmax(axis=0)  # [N]
+                    
+                    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
+                    
+                    top1_image_blended = []
+                    cm = matplotlib.colormaps['hot']
+                    for i_fps in range(len(top1_image_idxs)):
+                        image_idx = top1_image_idxs[i_fps]
+                        image = images[image_idx]
+                        heatmap = heatmaps[image_idx, :, :, i_fps]
+                        heatmap = cm(heatmap)
+                        heatmap = (heatmap[:, :, :3] * 255).astype(np.uint8)
+                        image = image.convert("RGB").resize((256, 256))
+                        heatmap = Image.fromarray(heatmap).resize((256, 256)).convert("RGB")
+                        blended = 0.5 * np.array(image) + 0.5 * np.array(heatmap)
+                        blended = np.clip(blended, 0, 255).astype(np.uint8)
+                        border_color = fps_tsne_rgb[i_fps, :3] * 255
+                        border_width = 20
+                        padded_image = pad_image_with_border(blended, border_color, border_width)
+                        top1_image_blended.append(padded_image)
+                    
+                    return top1_image_blended, dots_idx
+                        
+                    
+                
+                def plot_tsne_with_image_heatmaps(images, eigvecs, fps_eigvecs, tsne_embed, fps_tsne_rgb, max_display_dots=100):
+                    top1_image_blended, dots_idx = get_top1_heatmap_for_each_dot(images, eigvecs, fps_eigvecs, max_display_dots, fps_tsne_rgb, tsne_embed)
+                    
+                    # Plot the t-SNE points
+                    fig, ax = plt.subplots(1, 1, figsize=(15, 15))
+                    ax.scatter(tsne_embed[:, 0], tsne_embed[:, 1], s=20, c=fps_tsne_rgb)
+                    ax.set_xticks([])
+                    ax.set_yticks([])
+                    ax.axis('off')
+                    ax.set_xlim(tsne_embed[:, 0].min()*1.1, tsne_embed[:, 0].max()*1.1)
+                    ax.set_ylim(tsne_embed[:, 1].min()*1.1, tsne_embed[:, 1].max()*1.1)
+                    
+                    # Add the top1_image_blended to the scatter plot
+                    for i, (x, y) in enumerate(tsne_embed[dots_idx]):
+                        img = top1_image_blended[i]
+                        img = np.array(img)
+                        imgbox = OffsetImage(img, zoom=0.15)
+                        ab = AnnotationBbox(imgbox, (x, y), frameon=False)
+                        ax.add_artist(ab)
+                    ax.scatter(tsne_embed[:, 0], tsne_embed[:, 1], s=20, c=fps_tsne_rgb)
+                    
+                    # 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)
+                    return pil_image
+                    
+                
+                def run_fps_tsne_hierarchical(image_gallery, eigvecs, num_sample_fps, perplexity_tsne, tsne3d_rgb, seed=0, max_display_dots=300):
                     if len(eigvecs) == 0:
                         gr.Warning("Please run NCUT first.")
                         return
+                    images = [image[0] for image in image_gallery]
+                    if isinstance(images[0], str):
+                        images = [Image.open(image) for image in images]
+                        
                     eigvecs = torch.tensor(eigvecs)
-                    eigvecs = eigvecs.reshape(-1, eigvecs.shape[-1])
+                    _eigvecs = eigvecs.reshape(-1, eigvecs.shape[-1])
                     gr.Info("Running FPS, t-SNE, and Hierarchical Clustering...", 3)
                     from ncut_pytorch.ncut_pytorch import farthest_point_sampling
                     from sklearn.manifold import TSNE
@@ -2357,8 +2444,8 @@ with demo:
                     torch.manual_seed(seed)
                     np.random.seed(seed)
                     
-                    fps_idx = farthest_point_sampling(eigvecs, num_sample_fps)
-                    fps_eigvecs = eigvecs[fps_idx]
+                    fps_idx = farthest_point_sampling(_eigvecs, num_sample_fps)
+                    fps_eigvecs = _eigvecs[fps_idx]
                     fps_eigvecs = fps_eigvecs.numpy()
                     
                     tsne3d_rgb = tsne3d_rgb.reshape(-1, 3)
@@ -2371,19 +2458,29 @@ with demo:
                         metric='cosine',
                         random_state=seed,
                     ).fit_transform(fps_eigvecs)
+                    # normalize = [-1, 1]
+                    tsne_embed[:, 0] = (tsne_embed[:, 0] - tsne_embed[:, 0].min()) / (tsne_embed[:, 0].max() - tsne_embed[:, 0].min()) * 2 - 1
+                    tsne_embed[:, 1] = (tsne_embed[:, 1] - tsne_embed[:, 1].min()) / (tsne_embed[:, 1].max() - tsne_embed[:, 1].min()) * 2 - 1
                     
                     edges = build_tree(tsne_embed)
+                    # edges = build_tree(fps_eigvecs, dist='cosine')
+                    # edges = build_tree(fps_tsne3d_rgb)
                     
                     # Plot the t-SNE points
                     pil_image = plot_tsne_tree(tsne_embed, edges, fps_tsne3d_rgb, 0)
                     
-                    return tsne_embed, edges, fps_eigvecs, fps_tsne3d_rgb, fps_idx, pil_image
+                    # Plot the t-SNE points with image heatmaps
+                    big_pil_image = plot_tsne_with_image_heatmaps(images, eigvecs, fps_eigvecs, tsne_embed, fps_tsne3d_rgb, max_display_dots)
+                    
+                    return tsne_embed, edges, fps_eigvecs, fps_tsne3d_rgb, fps_idx, pil_image, big_pil_image
 
-                run_hierarchical_button.click(
-                    run_fps_tsne_hierarchical,
-                    inputs=[eigvecs, num_sample_fps_slider, tsne_perplexity_slider, tsne3d_rgb, fps_hc_seed_slider],
-                    outputs=[tsne_2d_points, edges, fps_eigvecs, fps_tsne_rgb, fps_indices, tsne_plot],
-                )
+        big_tsne_plot = gr.Image(label="spectral-tSNE tree [+ Cluster Heatmap]", elem_id="big_tsne_plot", interactive=False, format='png')
+        
+        run_hierarchical_button.click(
+            run_fps_tsne_hierarchical,
+            inputs=[input_gallery, eigvecs, num_sample_fps_slider, tsne_perplexity_slider, tsne3d_rgb, fps_hc_seed_slider],
+            outputs=[tsne_2d_points, edges, fps_eigvecs, fps_tsne_rgb, fps_indices, tsne_plot, big_tsne_plot],
+        )
         gr.Markdown('---')
         gr.Markdown('<h3 style="text-align: center;">↓ interactively inspect the hierarchical structure</h3>')
         gr.Markdown('---')