fix msg, faster eig norm

app.py

@@ -2232,6 +2232,12 @@ with demo:
                 def __run_fn(*args, **kwargs):
                     eigvecs, rgb, logging_str = run_fn(*args, **kwargs)
                     rgb_gallery = to_pil_images(rgb)
+                    # normalize the eigvecs
+                    eigvecs = torch.tensor(eigvecs)
+                    if torch.cuda.is_available():
+                        eigvecs = eigvecs.cuda()
+                    eigvecs = F.normalize(eigvecs, p=2, dim=-1)
+                    eigvecs = eigvecs.cpu().numpy()
                     return eigvecs, rgb, rgb_gallery, logging_str
                 
                 submit_button.click(
@@ -2408,8 +2414,8 @@ with demo:
                 gr.Markdown("Known Issue: Resize the browser window will break the clicking, please refresh the page.")
                 with gr.Accordion("Outputs", open=True):
                     gr.Markdown("""
-                                1. spectral-tSNE tree: ◆ means connected components to the selected point.
-                                2. Cluster Heatmap: max cosine similarity to any points in the connected components.
+                                1. spectral-tSNE tree: ◆ marker is the N points, connected components to the clicked dot.
+                                2. Cluster Heatmap: max of N cosine similarity to N points in the connected components.
                                 """)
             with gr.Column(scale=5, min_width=200): 
                 prompt_radio = gr.Radio(["Tree", "Image"], label="Where to click on?", value="Tree", elem_id="prompt_radio", show_label=True)
@@ -2427,6 +2433,7 @@ with demo:
                 tsne_plot.change(updaste_tsne_plot_change_granularity,
                                     inputs=[granularity_slider, tsne_2d_points, edges, fps_tsne_rgb],
                                     outputs=[tsne_prompt_image])
+                prompt_radio.change(update_image_prompt, inputs=[image_slider, output_gallery], outputs=[image_plot])
                 run_inspection_button = gr.Button("🔴 RUN Inspection", elem_id="run_inspection", variant='primary')
                 inspect_logging_text = gr.Textbox("Logging information", lines=3, label="Logging", elem_id="inspect_logging", type="text", placeholder="Logging information", autofocus=False, autoscroll=False)
                 # output_slot_radio = gr.Radio([1, 2, 3], label="Output Row", value=1, elem_id="output_slot", show_label=True)
@@ -2500,8 +2507,8 @@ with demo:
             x = int(x * w)
             y = int(y * h)
             eigvec = _eigvec[y, x]
-            dist = np.linalg.norm(fps_eigvecs - eigvec, axis=1)
-            closest_idx = np.argmin(dist)
+            sim = fps_eigvecs @ eigvec
+            closest_idx = np.argmax(sim)
             return closest_idx
         
         def find_closest_fps_point(prompt_radio, tsne_prompt, image_prompt, i_image, tsne2d_embed, eigvecs, fps_eigvecs):
@@ -2529,13 +2536,29 @@ with demo:
             output_tsne_plot = plot_tsne_tree(tsne2d_embed, edges, fps_tsne_rgb, granularity, closest_idx, highlight_connections=True)
             
             # draw heatmap for the connected components
+            ## cosine distance
             connected_eigvecs = fps_eigvecs[connected_idxs]
-            left = torch.tensor(eigvecs).float()  # B H W 3
-            right = torch.tensor(connected_eigvecs).float()
-            left = F.normalize(left, p=2, dim=-1)
-            right = F.normalize(right, p=2, dim=-1)
+            left = eigvecs.astype(np.float32)  # B H W C
+            right = connected_eigvecs.astype(np.float32)  # N C
+            # left = F.normalize(left, p=2, dim=-1)
+            # right = F.normalize(right, p=2, dim=-1)
+            # eigvec is already normalized when saved to gr.State
             similarity = left @ right.T
-            similarity = similarity.max(dim=-1).values  # B H W
+            similarity = similarity.max(axis=-1)  # B H W N
+            ## euclidean distance
+            # b, h, w = tsne3d_rgb.shape[:3]
+            # tsne3d_rgb = tsne3d_rgb.reshape(b*h*w, 3)
+            # connected_rgb = tsne3d_rgb[fps_indices][connected_idxs]
+            # left = torch.tensor(tsne3d_rgb).float()  # (B H W) 3
+            # right = torch.tensor(connected_rgb).float()  # N 3
+            # # dist B H W N
+            # dist = left[:, None] - right[None]
+            # dist = torch.sqrt((dist ** 2).sum(dim=-1))
+            # dist = dist.min(dim=-1).values  # B H W
+            # dist = dist.reshape(b, h, w)
+            # gr.Info(f"dist: min={dist.min().item()}, max={dist.max().item()}, mean={dist.mean().item()}", 3)
+            # similarity = 1 - dist
+            
             hot_map = matplotlib.colormaps['hot']
             heatmap = hot_map(similarity)[..., :3]  # B H W 3
             heatmap_images = to_pil_images(torch.tensor(heatmap), target_size=256, force_size=True)

fps_cluster.py

@@ -5,7 +5,8 @@ import torch
 
 def build_tree(all_dots):
     num_sample = all_dots.shape[0]
-    center = all_dots.mean(axis=0)
+    # center = all_dots.mean(axis=0)
+    center = np.median(all_dots, axis=0)
     distances_to_center = np.linalg.norm(all_dots - center, axis=1)
     start_idx = np.argmin(distances_to_center)
     indices = [start_idx]