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