Add sample image return functionality to app

pulls sample image from EOL subset of TreeOfLife-10M

app.py

@@ -6,12 +6,14 @@ import logging
 
 import gradio as gr
 import numpy as np
+import polars as pl
 import torch
 import torch.nn.functional as F
 from open_clip import create_model, get_tokenizer
 from torchvision import transforms
 
 from templates import openai_imagenet_template
+from components.query import  get_sample
 
 log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
 logging.basicConfig(level=logging.INFO, format=log_format)
@@ -19,6 +21,12 @@ logger = logging.getLogger()
 
 hf_token = os.getenv("HF_TOKEN")
 
+# For sample images
+METADATA_PATH = "components/metadata.csv"
+# Read page ID as int and filter out smaller ablation duplicated training split
+metadata_df = pl.read_csv(METADATA_PATH, low_memory = False)
+metadata_df = metadata_df.with_columns(pl.col("eol_page_id").cast(pl.Int64))
+
 model_str = "hf-hub:imageomics/bioclip"
 tokenizer_str = "ViT-B-16"
 
@@ -123,12 +131,14 @@ def format_name(taxon, common):
 
 
 @torch.no_grad()
-def open_domain_classification(img, rank: int) -> dict[str, float]:
+def open_domain_classification(img, rank: int, return_all=False):
     """
     Predicts from the entire tree of life.
     If targeting a higher rank than species, then this function predicts among all
     species, then sums up species-level probabilities for the given rank.
     """
+
+    logger.info(f"Starting open domain classification for rank: {rank}")
     img = preprocess_img(img).to(device)
     img_features = model.encode_image(img.unsqueeze(0))
     img_features = F.normalize(img_features, dim=-1)
@@ -136,21 +146,36 @@ def open_domain_classification(img, rank: int) -> dict[str, float]:
     logits = (model.logit_scale.exp() * img_features @ txt_emb).squeeze()
     probs = F.softmax(logits, dim=0)
 
-    # If predicting species, no need to sum probabilities.
     if rank + 1 == len(ranks):
         topk = probs.topk(k)
-        return {
+        prediction_dict = {
             format_name(*txt_names[i]): prob for i, prob in zip(topk.indices, topk.values)
         }
+        logger.info(f"Top K predictions: {prediction_dict}")
+        top_prediction_name = format_name(*txt_names[topk.indices[0]]).split("(")[0]
+        logger.info(f"Top prediction name: {top_prediction_name}")
+        sample_img, taxon_url = get_sample(metadata_df, top_prediction_name, rank)
+        if return_all:
+            return prediction_dict, sample_img, taxon_url
+        return prediction_dict
 
-    # Sum up by the rank
     output = collections.defaultdict(float)
     for i in torch.nonzero(probs > min_prob).squeeze():
         output[" ".join(txt_names[i][0][: rank + 1])] += probs[i]
 
     topk_names = heapq.nlargest(k, output, key=output.get)
+    prediction_dict = {name: output[name] for name in topk_names}
+    logger.info(f"Top K names for output: {topk_names}")
+    logger.info(f"Prediction dictionary: {prediction_dict}")
+
+    top_prediction_name = topk_names[0]
+    logger.info(f"Top prediction name: {top_prediction_name}")
+    sample_img, taxon_url = get_sample(metadata_df, top_prediction_name, rank)
+    logger.info(f"Sample image and taxon URL: {sample_img}, {taxon_url}")
 
-    return {name: output[name] for name in topk_names}
+    if return_all:
+        return prediction_dict, sample_img, taxon_url
+    return prediction_dict
 
 
 def change_output(choice):
@@ -179,9 +204,22 @@ if __name__ == "__main__":
         status_msg = f"{done}/{total} ({done / total * 100:.1f}%) indexed"
 
     with gr.Blocks() as app:
-        img_input = gr.Image()
-
+        
         with gr.Tab("Open-Ended"):
+            with gr.Row(variant = "panel", elem_id = "images_panel"):
+                with gr.Column():
+                    img_input = gr.Image(height = 400, sources=["upload"])
+            
+                with gr.Column():
+                    # display sample image of top predicted taxon
+                    sample_img = gr.Image(label = "Sample Image of Predicted Taxon", 
+                                        height = 400, 
+                                        show_download_button = False)
+                
+                    taxon_url = gr.HTML(label = "More Information", 
+                                    elem_id = "url"
+                                    )
+
             with gr.Row():
                 with gr.Column():
                     rank_dropdown = gr.Dropdown(
@@ -199,32 +237,36 @@ if __name__ == "__main__":
                         show_label=True,
                         value=None,
                     )
-                    # open_domain_flag_btn = gr.Button("Flag Mistake", variant="primary")
+                  #  open_domain_flag_btn = gr.Button("Flag Mistake", variant="primary")
 
             with gr.Row():
                 gr.Examples(
                     examples=open_domain_examples,
                     inputs=[img_input, rank_dropdown],
                     cache_examples=True,
-                    fn=open_domain_classification,
+                    fn=lambda img, rank: open_domain_classification(img, rank, return_all=False),
                     outputs=[open_domain_output],
                 )
-
-            # open_domain_callback = gr.HuggingFaceDatasetSaver(
-            #     hf_token, "imageomics/bioclip-demo-open-domain-mistakes", private=True
-            # )
-            # open_domain_callback.setup(
-            #     [img_input, rank_dropdown, open_domain_output],
-            #     flagging_dir="logs/flagged",
-            # )
-            # open_domain_flag_btn.click(
-            #     lambda *args: open_domain_callback.flag(args),
-            #     [img_input, rank_dropdown, open_domain_output],
-            #     None,
-            #     preprocess=False,
-            # )
-
+            '''
+            # Flagging Code
+            open_domain_callback = gr.HuggingFaceDatasetSaver(
+                hf_token, "bioclip-demo-open-domain-mistakes", private=True
+            )
+            open_domain_callback.setup(
+                [img_input, rank_dropdown, open_domain_output],
+                flagging_dir="bioclip-demo-open-domain-mistakes/logs/flagged",
+            )
+            open_domain_flag_btn.click(
+                lambda *args: open_domain_callback.flag(args),
+                [img_input, rank_dropdown, open_domain_output],
+                None,
+                preprocess=False,
+            )
+            '''
         with gr.Tab("Zero-Shot"):
+            with gr.Row():
+                img_input_zs = gr.Image(height = 400, sources=["upload"])
+
             with gr.Row():
                 with gr.Column():
                     classes_txt = gr.Textbox(
@@ -240,43 +282,44 @@ if __name__ == "__main__":
                     zero_shot_output = gr.Label(
                         num_top_classes=k, label="Prediction", show_label=True
                     )
-                    # zero_shot_flag_btn = gr.Button("Flag Mistake", variant="primary")
+             #       zero_shot_flag_btn = gr.Button("Flag Mistake", variant="primary")
 
             with gr.Row():
                 gr.Examples(
                     examples=zero_shot_examples,
-                    inputs=[img_input, classes_txt],
+                    inputs=[img_input_zs, classes_txt],
                     cache_examples=True,
                     fn=zero_shot_classification,
                     outputs=[zero_shot_output],
                 )
-
-        # zero_shot_callback = gr.HuggingFaceDatasetSaver(
-        #     hf_token, "imageomics/bioclip-demo-zero-shot-mistakes", private=True
-        # )
-        # zero_shot_callback.setup(
-        #     [img_input, zero_shot_output], flagging_dir="logs/flagged"
-        # )
-        # zero_shot_flag_btn.click(
-        #     lambda *args: zero_shot_callback.flag(args),
-        #     [img_input, zero_shot_output],
-        #     None,
-        #     preprocess=False,
-        # )
-
+        '''
+        # Flagging Code
+        zero_shot_callback = gr.HuggingFaceDatasetSaver(
+            hf_token, "bioclip-demo-zero-shot-mistakes", private=True
+        )
+        zero_shot_callback.setup(
+            [img_input, zero_shot_output], flagging_dir="bioclip-demo-zero-shot-mistakes/logs/flagged"
+        )
+        zero_shot_flag_btn.click(
+            lambda *args: zero_shot_callback.flag(args),
+            [img_input, zero_shot_output],
+            None,
+            preprocess=False,
+        )
+        '''
         rank_dropdown.change(
             fn=change_output, inputs=rank_dropdown, outputs=[open_domain_output]
         )
 
         open_domain_btn.click(
-            fn=open_domain_classification,
+            fn=lambda img, rank: open_domain_classification(img, rank, return_all=True),
             inputs=[img_input, rank_dropdown],
-            outputs=[open_domain_output],
+            outputs=[open_domain_output, sample_img, taxon_url],
         )
 
         zero_shot_btn.click(
             fn=zero_shot_classification,
-            inputs=[img_input, classes_txt],
+            inputs=[img_input_zs, classes_txt],
             outputs=zero_shot_output,
         )
         
@@ -291,4 +334,4 @@ if __name__ == "__main__":
         )
 
     app.queue(max_size=20)
-    app.launch()
+    app.launch(share=True)