Closest neighbors: start at 5 and allow 'Add K more neighbors' with dynamic K and tables

app.py

@@ -76,7 +76,7 @@ USER_GUIDE = """
 
 - **Input**: Upload a fossil-leaf photograph (ideally a tightly cropped Florissant compression fossil with minimal background clutter).\n
 - **Output A — predicted families**: The model suggests the top-ranked plant families (as hypotheses to guide expert review).\n
-- **Output B — closest training fossils**: The most similar Florissant fossils in the reference set (for qualitative comparison).\n
+- **Output B — closest training fossils**: The most similar Florissant fossils in [The image dataset](https://plus.figshare.com/articles/dataset/Image_collection_and_supporting_data_for_An_image_dataset_of_cleared_x-rayed_and_fossil_leaves_vetted_to_plant_family_for_human_and_machine_learning/14980698) (for qualitative comparison).\n
 - **Output C — explanations**: Heatmaps showing image regions that support each suggested family (qualitative).\n
 
 ### How to use
@@ -251,17 +251,17 @@ def _top_predicted_class(class_predicted):
     return None
 
 
-def find_closest(input_image, model_name, predicted_class=None):
+def find_closest(input_image, model_name, predicted_class=None, top_k=5):
     embedding = get_embeddings(input_image, model_name)
     from closest_sample import get_images, get_images_fossils
-    classes, paths, filenames = get_images(embedding, model_name, predicted_class=predicted_class)
+    classes, paths, filenames = get_images(embedding, model_name, predicted_class=predicted_class, top_k=top_k)
     return classes, paths, filenames
 
 
-def find_closest_fossils(input_image, model_name, predicted_class=None):
+def find_closest_fossils(input_image, model_name, predicted_class=None, top_k=5):
     embedding = get_embeddings(input_image, model_name)
     from closest_sample import get_images_fossils
-    return get_images_fossils(embedding, model_name, predicted_class=predicted_class)
+    return get_images_fossils(embedding, model_name, predicted_class=predicted_class, top_k=top_k)
 
 def generate_diagram_closest(input_image,model_name,top_k):
     embedding = get_embeddings(input_image,model_name)
@@ -487,17 +487,17 @@ def setup_examples():
     # Gradio Examples can handle URLs directly - they will fetch and display the images
     # Pass URLs as the first argument - Gradio will automatically fetch and display them
     # Note: Gradio downloads URLs to temp directory, which is normal behavior
-    print(f"DEBUG: Final fossil_samples count: {len(fossil_samples)}")
+        print(f"DEBUG: Final fossil_samples count: {len(fossil_samples)}")
     if len(fossil_samples) > 0:
         print(f"DEBUG: First fossil sample (should be URL): {fossil_samples[0]}")
         print(f"DEBUG: Is URL: {fossil_samples[0].startswith('http') if fossil_samples else False}")
     
     examples_fossils = gr.Examples(
-        fossil_samples, 
+        fossil_samples,
         inputs=input_image,
         examples_per_page=6,  # Reduced for better spacing and organization
-        label='Leaf fossil examples from the dataset',
-        elem_id="fossil-examples"
+        label='Leaf fossil examples from The image dataset',
+        elem_id="fossil-examples",
     )
     return examples_fossils
 
@@ -842,7 +842,9 @@ with gr.Blocks(theme='sudeepshouche/minimalist', css=custom_css) as demo:
                     with gr.Accordion("📸 Browse Florissant fossils (non-NA)", open=True):
                         gr.Markdown(
                             "<p style='font-size: 13px; margin-bottom: 10px;'>"
-                            "These thumbnails are sourced from the Florissant dataset, from specimens where there are doubts about their family. "
+                                "These thumbnails are sourced from the Florissant subset of "
+                                "<a href='https://plus.figshare.com/articles/dataset/Image_collection_and_supporting_data_for_An_image_dataset_of_cleared_x-rayed_and_fossil_leaves_vetted_to_plant_family_for_human_and_machine_learning/14980698' target='_blank'>The image dataset</a>, "
+                                "from specimens where there are doubts about their family. "
                             "For full context pages, use: "
                             "<a href='https://serre-lab.github.io/FossilLeafLens/' target='_blank'>Fossil Leaf Lens</a>."
                             "</p>"
@@ -971,20 +973,28 @@ with gr.Blocks(theme='sudeepshouche/minimalist', css=custom_css) as demo:
                     #             label_closest_image_4 = gr.Markdown('')
                     #             closest_image_4 = gr.Image(label='Fifth Closest Image',image_mode='contain',width=200, height=200)
                     #     find_closest_btn = gr.Button("Find Closest Images")
+                    neighbors_k_state = gr.State(5)
+                    neighbors_add_k = gr.Slider(1, 20, value=5, step=1, label="Add K neighbors", info="Number of additional neighbors to fetch on demand")
+
                     with gr.Accordion('Closest Extant Leaves'):
                         gr.Markdown(
-                            "5 closest extant leaves (2024 dataset). **Ordered from most to least similar.** "
+                            "5 closest extant leaves from the extant subset of "
+                            "<a href='https://plus.figshare.com/articles/dataset/Image_collection_and_supporting_data_for_An_image_dataset_of_cleared_x-rayed_and_fossil_leaves_vetted_to_plant_family_for_human_and_machine_learning/14980698' target='_blank'>The image dataset</a>. "
+                            "**Ordered from most to least similar.** "
                             "Neighbors are chosen by *embedding* similarity (visual), not by predicted family; "
                             "if you ran **Classify Image** first, we try to include at least one from the top-predicted family when possible."
                         )
                         closest_table = gr.HTML(label="Closest Extant Table")
                     with gr.Accordion('Closest Fossils'):
                         gr.Markdown(
-                            "5 closest fossil images (reference fossil dataset). **Ordered from most to least similar.** "
-                            "Same logic: embedding similarity + at least one from predicted family when you have run Classify Image."
+                            "5 closest Florissant fossil images from the fossil subset of "
+                            "<a href='https://plus.figshare.com/articles/dataset/Image_collection_and_supporting_data_for_An_image_dataset_of_cleared_x-rayed_and_fossil_leaves_vetted_to_plant_family_for_human_and_machine_learning/14980698' target='_blank'>The image dataset</a>. "
+                            "**Ordered from most to least similar.** "
+                            "Same logic: embedding similarity + at least one from predicted family when you have run **Classify Image**."
                         )
                         closest_fossils_table = gr.HTML(label="Closest Fossils Table")
                     find_closest_btn = gr.Button("Find Closest Images", icon="https://www.svgrepo.com/show/13672/play-button.svg")
+                    add_neighbors_btn = gr.Button("Add K more neighbors", icon="https://www.svgrepo.com/show/13672/play-button.svg")
                     
                     #segment_button.click(segment_image, inputs=input_image, outputs=segmented_image)
                     classify_image_button.click(classify_image, inputs=[original_image,model_name], outputs=class_predicted)
@@ -992,7 +1002,10 @@ with gr.Blocks(theme='sudeepshouche/minimalist', css=custom_css) as demo:
                     # with gr.Accordion('Closest Leaves Images'):
                     #     gr.Markdown("5 closest leaves")
                     with gr.Accordion("Family Distribution of Closest Samples "):
-                        gr.Markdown("Visualize plant family distribution of top-k closest samples in our dataset")
+                        gr.Markdown(
+                            "Visualize plant family distribution of top-k closest samples in "
+                            "<a href='https://plus.figshare.com/articles/dataset/Image_collection_and_supporting_data_for_An_image_dataset_of_cleared_x-rayed_and_fossil_leaves_vetted_to_plant_family_for_human_and_machine_learning/14980698' target='_blank'>The image dataset</a>."
+                        )
                         with gr.Column():
                             with gr.Row():
                                 diagram= gr.Image(label = 'Bar Chart')
@@ -1042,7 +1055,8 @@ with gr.Blocks(theme='sudeepshouche/minimalist', css=custom_css) as demo:
                 <thead><tr><th>Rank (1 = most similar)</th><th>Image</th><th>Plant Family</th><th>Specimen Name</th></tr></thead>
                 <tbody>
             """
-            for i in range(5):
+            n = min(len(images), len(labels), len(filenames)) if filenames is not None else min(len(images), len(labels))
+            for i in range(n):
                 rank = i + 1
                 img_src = ""
                 if i < len(images):
@@ -1077,13 +1091,34 @@ with gr.Blocks(theme='sudeepshouche/minimalist', css=custom_css) as demo:
             table_html += "</tbody></table>"
             return table_html
 
-        def update_closest_outputs(input_image, model_name, class_predicted):
+        def _compute_closest(input_image, model_name, class_predicted, k):
             predicted_class = _top_predicted_class(class_predicted)
-            extant_labels, extant_images, extant_filenames = find_closest(input_image, model_name, predicted_class=predicted_class)
-            fossil_labels, fossil_images, fossil_filenames = find_closest_fossils(input_image, model_name, predicted_class=predicted_class)
-            return _closest_table_html(extant_labels, extant_images, extant_filenames), _closest_table_html(fossil_labels, fossil_images, fossil_filenames)
-
-        find_closest_btn.click(fn=update_closest_outputs, inputs=[original_image, model_name, class_predicted], outputs=[closest_table, closest_fossils_table])
+            extant_labels, extant_images, extant_filenames = find_closest(input_image, model_name, predicted_class=predicted_class, top_k=k)
+            fossil_labels, fossil_images, fossil_filenames = find_closest_fossils(input_image, model_name, predicted_class=predicted_class, top_k=k)
+            return (
+                _closest_table_html(extant_labels, extant_images, extant_filenames),
+                _closest_table_html(fossil_labels, fossil_images, fossil_filenames),
+                k,
+            )
+
+        def update_closest_outputs_initial(input_image, model_name, class_predicted):
+            # Reset to 5 neighbors
+            return _compute_closest(input_image, model_name, class_predicted, k=5)
+
+        def update_closest_outputs_add(input_image, model_name, class_predicted, current_k, add_k):
+            new_k = int(current_k) + int(add_k)
+            return _compute_closest(input_image, model_name, class_predicted, k=new_k)
+
+        find_closest_btn.click(
+            fn=update_closest_outputs_initial,
+            inputs=[original_image, model_name, class_predicted],
+            outputs=[closest_table, closest_fossils_table, neighbors_k_state],
+        )
+        add_neighbors_btn.click(
+            fn=update_closest_outputs_add,
+            inputs=[original_image, model_name, class_predicted, neighbors_k_state, neighbors_add_k],
+            outputs=[closest_table, closest_fossils_table, neighbors_k_state],
+        )
         #classify_segmented_button.click(classify_image, inputs=[segmented_image,model_name], outputs=class_predicted)
 
         generate_diagram.click(generate_diagram_closest, inputs=[original_image,model_name,top_k], outputs=diagram)

closest_sample.py

@@ -277,7 +277,7 @@ def download_public_image(url, destination_path):
     else:
         print(f"Failed to download image from bucket. Status code: {response.status_code}")
 
-def get_images(embedding, model_name, predicted_class=None):
+def get_images(embedding, model_name, predicted_class=None, top_k=5):
     if model_name in ['Rock 170', 'Mummified 170']:
         pca_fossils = load_pickle_safe('pca_fossils_170_finer.pkl')
         pca_leaves = load_pickle_safe('pca_leaves_170_finer.pkl')
@@ -296,10 +296,10 @@ def get_images(embedding, model_name, predicted_class=None):
                     families.append(parts[-2] if len(parts) >= 2 else "Unknown")
                 if predicted_class:
                     pca_d = pca_distance_simple_with_predicted(
-                        pca_fossils, embedding, embedding_fossils_2024, 5, families, predicted_class
+                        pca_fossils, embedding, embedding_fossils_2024, top_k, families, predicted_class
                     )
                 else:
-                    pca_d = pca_distance_simple(pca_fossils, embedding, embedding_fossils_2024, top_k=5)
+                    pca_d = pca_distance_simple(pca_fossils, embedding, embedding_fossils_2024, top_k=top_k)
                 local_paths = []
                 classes = []
                 filenames = []
@@ -322,7 +322,7 @@ def get_images(embedding, model_name, predicted_class=None):
         raise ValueError(f'{model_name} not recognized')
     #pca_embedding_fossils = pca_fossils.transform(embedding_fossils[:,-1])
     
-    pca_d = pca_distance(pca_fossils, embedding, embedding_fossils, top_k=5)
+    pca_d = pca_distance(pca_fossils, embedding, embedding_fossils, top_k=top_k)
     
     fossils_paths = fossils_pd['file_name'].values
     
@@ -374,7 +374,7 @@ def get_images(embedding, model_name, predicted_class=None):
     return classes, local_paths, filenames
 
 
-def get_images_fossils(embedding, model_name, predicted_class=None):
+def get_images_fossils(embedding, model_name, predicted_class=None, top_k=5):
     """
     Use 2024 fossil embeddings (from fossils path) when available;
     otherwise fall back to original finer.npy + fossils_pd + GCS. Fossils 142 / BEiT only.
@@ -396,10 +396,10 @@ def get_images_fossils(embedding, model_name, predicted_class=None):
                 families.append(parts[-2] if len(parts) >= 2 else "Unknown")
             if predicted_class:
                 pca_d = pca_distance_simple_with_predicted(
-                    pca_fossils, embedding, embedding_fossils_2024, 5, families, predicted_class
+                    pca_fossils, embedding, embedding_fossils_2024, top_k, families, predicted_class
                 )
             else:
-                pca_d = pca_distance_simple(pca_fossils, embedding, embedding_fossils_2024, top_k=5)
+                pca_d = pca_distance_simple(pca_fossils, embedding, embedding_fossils_2024, top_k=top_k)
             local_paths = []
             classes = []
             filenames = []