Sample image return (#8)

- Add a to learn more URL at the bottom of the app linking to GitHub and TreeOfLife dataset (6f2766a9fbe1c26c96934dfe41049e05f7ae2cae)
- update requirements for sample image return (9e5fdea1a9b91ad590f961bf1a2cc028d4a6d161)
- Add components for sample image return (1ca4d249a8b99b33200cf026163194de98e167cd)
- Add sample image return functionality to app (16df1e71bfc77b4a496b760a2161d3242177f62a)

.gitattributes

@@ -33,7 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-
 *.json filter=lfs diff=lfs merge=lfs -text
 *.jpeg filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
+components/metadata.csv filter=lfs diff=lfs merge=lfs -text

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,45 +282,56 @@ 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,
         )
+        
+        # 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()
+    app.launch(share=True)

components/metadata.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8576f6ca106f35387506369a70df01fb92192a740c3b5da2a12ad8303976aad
+size 233934143

components/metadata_readme.md

@@ -0,0 +1,11 @@
+---
+title: Bioclip Demo
+emoji: 🐘
+colorFrom: indigo
+colorTo: purple
+sdk: gradio
+sdk_version: 4.36.1
+app_file: app.py
+pinned: false
+license: mit
+---

components/query.py

@@ -0,0 +1,116 @@
+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))
+        full_eol_url = EOL_URL + eol_page_id
+        if is_exact:
+            eol_page = f"<p>Check out the EOL entry for {pred_taxon} to learn more: <a href={full_eol_url} target='_blank'>{full_eol_url}</a>.</p>"
+        else:
+            eol_page = f"<p>Check out an example EOL entry within {pred_taxon} to learn more: {full_name} <a href={full_eol_url} target='_blank'>{full_eol_url}</a>.</p>"
+        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

components/sync_samples_to_s3.bash

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+<<COMMENT
+Usage: 
+bash sync_samples_to_s3.bash <BASE_DIR>
+
+Dependencies: 
+- awscli (https://aws.amazon.com/cli/)
+Credentials to export as environment variables: 
+- AWS_ACCESS_KEY_ID
+- AWS_SECRET_ACCESS_KEY
+COMMENT
+
+# Check if a valid directory is provided as an argument
+if [ -z "$1" ]; then
+  echo "Usage: $0 <BASE_DIR>"
+  exit 1
+fi
+
+if [ ! -d "$1" ]; then
+  echo "Error: $1 is not a valid directory"
+  exit 1
+fi
+
+BASE_DIR="$1"
+S3_BUCKET="s3://treeoflife-10m-sample-images"
+
+# Loop through all directories and sync them to S3
+for dir in $BASE_DIR/*; do
+  if [ -d "$dir" ]; then
+    dir_name=$(basename "$dir")
+    aws s3 sync "$dir" "$S3_BUCKET/$dir_name/"
+  fi
+done

requirements.txt

@@ -1,4 +1,7 @@
 open_clip_torch
 torchvision
 torch
-gradio
+gradio
+polars
+pillow
+boto3