Retrieve example TOL-10M image and representative EOL page for OE prediction

Co-authored by: Elizabeth Campolongo <egrace479@users.noreply.huggingface.co>

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}")
 
-    return {name: output[name] for name in topk_names}
+    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}")
+
+    if return_all:
+        return prediction_dict, sample_img, taxon_url
+    return prediction_dict
 
 
 def change_output(choice):
@@ -179,9 +204,19 @@ 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.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)
+
             with gr.Row():
                 with gr.Column():
                     rank_dropdown = gr.Dropdown(
@@ -201,12 +236,17 @@ if __name__ == "__main__":
                     )
                   #  open_domain_flag_btn = gr.Button("Flag Mistake", variant="primary")
 
+            with gr.Row():
+                taxon_url = gr.TextArea(label = "More Information", 
+                                elem_id = "textbox", 
+                                show_copy_button = True)
+
             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],
                 )
             '''
@@ -225,6 +265,9 @@ if __name__ == "__main__":
             )
             '''
         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(
@@ -245,7 +288,7 @@ if __name__ == "__main__":
             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],
@@ -268,17 +311,33 @@ if __name__ == "__main__":
             fn=change_output, inputs=rank_dropdown, outputs=[open_domain_output]
         )
 
+        # open_domain_btn.click(
+        #     fn=open_domain_classification,
+        #     inputs=[img_input, rank_dropdown],
+        #     outputs=[open_domain_output, sample_img, taxon_url],
+        # )
+
         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,
         )
+        
+        # Footer to point out to model and data from app page.
+        gr.Markdown(
+            """
+            For more information on the [BioCLIP Model](https://huggingface.co/imageomics/bioclip) creation, see our [BioCLIP Project GitHub](https://github.com/Imageomics/bioclip), and
+            for easier integration of BioCLIP, checkout [pybioclip](https://github.com/Imageomics/pybioclip).
+            
+            To learn more about the data, check out our [TreeOfLife-10M Dataset](https://huggingface.co/datasets/imageomics/TreeOfLife-10M).
+            """
+        )
 
     app.queue(max_size=20)
     app.launch()

components/query.py

@@ -0,0 +1,115 @@
+import io
+import boto3
+import requests
+import numpy as np
+import polars as pl
+from PIL import Image
+from botocore.config import Config
+import logging
+
+logger = logging.getLogger(__name__)
+
+# S3 for sample images
+my_config = Config(
+    region_name='us-east-1'
+)
+s3_client = boto3.client('s3', config=my_config)
+
+# Set basepath for EOL pages for info
+EOL_URL = "https://eol.org/pages/"
+RANKS = ["kingdom", "phylum", "class", "order", "family", "genus", "species"]
+
+def get_sample(df, pred_taxon, rank):
+    '''
+    Function to retrieve a sample image of the predicted taxon and EOL page link for more info.
+
+    Parameters:
+    -----------
+    df : DataFrame
+        DataFrame with all sample images listed and their filepaths (in "file_path" column).
+    pred_taxon : str
+        Predicted taxon of the uploaded image.
+    rank : int
+        Index of rank in RANKS chosen for prediction.
+
+    Returns:
+    --------
+    img : PIL.Image
+        Sample image of predicted taxon for display.
+    eol_page : str
+        URL to EOL page for the taxon (may be a lower rank, e.g., species sample).
+    '''
+    logger.info(f"Getting sample for taxon: {pred_taxon} at rank: {rank}")
+    try:
+        filepath, eol_page_id, full_name, is_exact = get_sample_data(df, pred_taxon, rank)
+    except Exception as e:
+        logger.error(f"Error retrieving sample data: {e}")
+        return None, f"We encountered the following error trying to retrieve a sample image: {e}."
+    if filepath is None:
+        logger.warning(f"No sample image found for taxon: {pred_taxon}")
+        return None, f"Sorry, our EOL images do not include {pred_taxon}."
+
+    # Get sample image of selected individual
+    try:
+        img_src = s3_client.generate_presigned_url('get_object',
+                                                   Params={'Bucket': 'treeoflife-10m-sample-images',
+                                                           'Key': filepath}
+                                                   )
+        img_resp = requests.get(img_src)
+        img = Image.open(io.BytesIO(img_resp.content))
+        if is_exact:
+            eol_page = f"Check out the EOL entry for {pred_taxon} to learn more: {EOL_URL}{eol_page_id}."
+        else:
+            eol_page = f"Check out an example EOL entry within {pred_taxon} to learn more: {full_name} {EOL_URL}{eol_page_id}."
+        logger.info(f"Successfully retrieved sample image and EOL page for {pred_taxon}")
+        return img, eol_page
+    except Exception as e:
+        logger.error(f"Error retrieving sample image: {e}")
+        return None, f"We encountered the following error trying to retrieve a sample image: {e}."
+
+def get_sample_data(df, pred_taxon, rank):
+    '''
+    Function to randomly select a sample individual of the given taxon and provide associated native location.
+    
+    Parameters:
+    -----------
+    df : DataFrame
+        DataFrame with all sample images listed and their filepaths (in "file_path" column).
+    pred_taxon : str
+        Predicted taxon of the uploaded image.
+    rank : int
+        Index of rank in RANKS chosen for prediction.
+
+    Returns:
+    --------
+    filepath : str
+        Filepath of selected sample image for predicted taxon.
+    eol_page_id : str
+        EOL page ID associated with predicted taxon for more information.
+    full_name : str
+        Full taxonomic name of the selected sample.
+    is_exact : bool
+        Flag indicating if the match is exact (i.e., with empty lower ranks).
+    '''
+    for idx in range(rank + 1):
+        taxon = RANKS[idx]
+        target_taxon = pred_taxon.split(" ")[idx]
+        df = df.filter(pl.col(taxon) == target_taxon)
+
+    if df.shape[0] == 0:
+        return None, np.nan, "", False
+
+    # First, try to find entries with empty lower ranks
+    exact_df = df
+    for lower_rank in RANKS[rank + 1:]:
+        exact_df = exact_df.filter((pl.col(lower_rank).is_null()) | (pl.col(lower_rank) == ""))
+
+    if exact_df.shape[0] > 0:
+        df_filtered = exact_df.sample()
+        full_name = " ".join(df_filtered.select(RANKS[:rank+1]).row(0))
+        return df_filtered["file_path"][0], df_filtered["eol_page_id"].cast(pl.String)[0], full_name, True
+
+    # If no exact matches, return any entry with the specified rank
+    df_filtered = df.sample()
+    full_name = " ".join(df_filtered.select(RANKS[:rank+1]).row(0)) + " " + " ".join(df_filtered.select(RANKS[rank+1:]).row(0))
+    return df_filtered["file_path"][0], df_filtered["eol_page_id"].cast(pl.String)[0], full_name, False

requirements.txt

@@ -2,3 +2,6 @@ open_clip_torch
 torchvision
 torch
 gradio
+polars
+pillow
+boto3