Set up BioCLIP 2 app code based on BioCLIP demo

.gitattributes

@@ -33,3 +33,12 @@ 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
+examples/Actinostola-abyssorum.png filter=lfs diff=lfs merge=lfs -text
+examples/Amanita-muscaria.jpeg filter=lfs diff=lfs merge=lfs -text
+examples/Carnegiea-gigantea.png filter=lfs diff=lfs merge=lfs -text
+examples/Felis-catus.jpeg filter=lfs diff=lfs merge=lfs -text
+examples/milk-snake.png filter=lfs diff=lfs merge=lfs -text
+examples/Onoclea-sensibilis.jpg filter=lfs diff=lfs merge=lfs -text
+examples/Phoca-vitulina.png filter=lfs diff=lfs merge=lfs -text
+examples/Sarcoscypha-coccinea.jpeg filter=lfs diff=lfs merge=lfs -text
+examples/Ursus-arctos.jpeg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -8,6 +8,27 @@ sdk_version: 5.33.0
 app_file: app.py
 pinned: false
 license: mit
+models:
+- imageomics/bioclip-2
+datasets:
+- imageomics/TreeOfLife-200M
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+This app is modified from the original [BioCLIP Demo](https://huggingface.co/spaces/imageomics/bioclip-demo) to run inference with [BioCLIP 2](https://huggingface.co/imageomics/bioclip-2) and uses [pybioclip](https://github.com/Imageomics/pybioclip).
+
+Due to space persistent storage limitations, embeddings are fetched from the [TreeOfLife-200M repo](https://huggingface.co/datasets/imageomics/TreeOfLife-200M). The images will be retrieved from an S3 bucket, as with the origin, described below.
+
+Note that if this space is duplicated, the sample image portion **will not work**.
+
+**bioclip-2/metadata.parquet:** metadata file for fetching [TreeOfLife-200M](https://huggingface.co/datasets/imageomics/TreeOfLife-200M) sample images (up to 3 available per taxa) from an S3 bucket.
+- `uuid`: unique identifier for the image within the TreeOfLife-200M dataset.
+- `eol_page_id`: identifier of EOL page for the most specific taxa of the image (where available). Note that an image's association to a particular page ID may change with updates to the EOL (or image provider's) hierarchy. However, EOL taxon page IDs are stable. "https://eol.org/pages/" + `eol_page_id` links to the page.
+- `gbif_id`: identifier used by GBIF for the most specific taxa of the image (where available). "https://gbif.org/species/" + `taxonID` links to the page.
+- `kingdom`: kingdom to which the subject of the image belongs (all `Animalia`).
+- `phylum`: phylum to which the subject of the image belongs.
+- `class`: class to which the subject of the image belongs.
+- `order`: order to which the subject of the image belongs.
+- `family`: family to which the subject of the image belongs.
+- `genus`: genus to which the subject of the image belongs.
+- `species`: species to which the subject of the image belongs.
+- `file_path`: image filepath to fetch image from S3 bucket (`<folder>/<uuid>.jpg`, folders are first two characters of the `uuid`).

app.py

@@ -0,0 +1,287 @@
+import collections
+import heapq
+import json
+import os
+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 components.query import  get_sample
+from bioclip import CustomLabelsClassifier
+
+log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
+logging.basicConfig(level=logging.INFO, format=log_format)
+logger = logging.getLogger()
+
+hf_token = os.getenv("HF_TOKEN")
+
+# For sample images
+METADATA_PATH = "components/metadata.parquet"
+# Read page IDs as int
+metadata_df = pl.read_parquet(METADATA_PATH, low_memory = False)
+metadata_df = metadata_df.with_columns(pl.col(["eol_page_id", "gbif_id"]).cast(pl.Int64))
+
+MODEL_STR = "hf-hub:imageomics/bioclip-2"
+TOKENIZER_STR = "ViT-L-14"
+
+txt_emb_npy = "https://huggingface.co/datasets/imageomics/TreeOfLife-200M/resolve/main/embeddings/txt_emb_species.npy"
+txt_names_json = "embeddings/txt_emb_species.json"
+
+min_prob = 1e-9
+k = 5
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+preprocess_img = transforms.Compose(
+    [
+        transforms.ToTensor(),
+        transforms.Resize((224, 224), antialias=True),
+        transforms.Normalize(
+            mean=(0.48145466, 0.4578275, 0.40821073),
+            std=(0.26862954, 0.26130258, 0.27577711),
+        ),
+    ]
+)
+
+ranks = ("Kingdom", "Phylum", "Class", "Order", "Family", "Genus", "Species")
+
+open_domain_examples = [
+    ["examples/Ursus-arctos.jpeg", "Species"],
+    ["examples/Phoca-vitulina.png", "Species"],
+    ["examples/Felis-catus.jpeg", "Genus"],
+    ["examples/Sarcoscypha-coccinea.jpeg", "Order"],
+]
+zero_shot_examples = [
+    [
+        "examples/Ursus-arctos.jpeg",
+        "brown bear\nblack bear\npolar bear\nkoala bear\ngrizzly bear",
+    ],
+    ["examples/milk-snake.png", "coral snake\nmilk snake"],
+    ["examples/coral-snake.jpeg", "coral snake\nmilk snake"],
+    [
+        "examples/Carnegiea-gigantea.png",
+        "Carnegiea gigantea\nSchlumbergera opuntioides\nMammillaria albicoma",
+    ],
+    [
+        "examples/Amanita-muscaria.jpeg",
+        "Amanita fulva\nAmanita vaginata (grisette)\nAmanita calyptrata (coccoli)\nAmanita crocea\nAmanita rubescens (blusher)\nAmanita caesarea (Caesar's mushroom)\nAmanita jacksonii (American Caesar's mushroom)\nAmanita muscaria (fly agaric)\nAmanita pantherina (panther cap)",
+    ],
+    [
+        "examples/Actinostola-abyssorum.png",
+        "Animalia Cnidaria Hexacorallia Actiniaria Actinostolidae Actinostola abyssorum\nAnimalia Cnidaria Hexacorallia Actiniaria Actinostolidae Actinostola bulbosa\nAnimalia Cnidaria Hexacorallia Actiniaria Actinostolidae Actinostola callosa\nAnimalia Cnidaria Hexacorallia Actiniaria Actinostolidae Actinostola capensis\nAnimalia Cnidaria Hexacorallia Actiniaria Actinostolidae Actinostola carlgreni",
+    ],
+    [
+        "examples/Sarcoscypha-coccinea.jpeg",
+        "scarlet elf cup (coccinea)\nscharlachroter kelchbecherling (austriaca)\ncrimson cup (dudleyi)\nstalked scarlet cup (occidentalis)",
+    ],
+    [
+        "examples/Onoclea-hintonii.jpg",
+        "Onoclea attenuata\nOnoclea boryana\nOnoclea hintonii\nOnoclea intermedia\nOnoclea sensibilis",
+    ],
+    [
+        "examples/Onoclea-sensibilis.jpg",
+        "Onoclea attenuata\nOnoclea boryana\nOnoclea hintonii\nOnoclea intermedia\nOnoclea sensibilis",
+    ],
+]
+
+
+def indexed(lst, indices):
+    return [lst[i] for i in indices]
+
+
+def zero_shot_classification(img, cls_str: str) -> dict[str, float]:
+    classes = [cls.strip() for cls in cls_str.split("\n") if cls.strip()]
+    classifier = CustomLabelsClassifier(
+        cls_ary = classes,
+        model_str = MODEL_STR,  # remove this line once pybioclip uses BioCLIP 2
+    )
+    return classifier.predict(img)
+
+
+def format_name(taxon, common):
+    taxon = " ".join(taxon)
+    if not common:
+        return taxon
+    return f"{taxon} ({common})"
+
+
+@torch.no_grad()
+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)
+
+    logits = (model.logit_scale.exp() * img_features @ txt_emb).squeeze()
+    probs = F.softmax(logits, dim=0)
+
+    if rank + 1 == len(ranks):
+        topk = probs.topk(k)
+        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
+
+    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}")
+
+    if return_all:
+        return prediction_dict, sample_img, taxon_url
+    return prediction_dict
+
+
+def change_output(choice):
+    return gr.Label(num_top_classes=k, label=ranks[choice], show_label=True, value=None)
+
+
+if __name__ == "__main__":
+    logger.info("Starting.")
+    model = create_model(MODEL_STR, output_dict=True, require_pretrained=True)
+    model = model.to(device)
+    logger.info("Created model.")
+
+    model = torch.compile(model)
+    logger.info("Compiled model.")
+
+    tokenizer = get_tokenizer(TOKENIZER_STR)
+
+    txt_emb = torch.from_numpy(np.load(txt_emb_npy, mmap_mode="r")).to(device)
+    with open(txt_names_json) as fd:
+        txt_names = json.load(fd)
+
+    done = txt_emb.any(axis=0).sum().item()
+    total = txt_emb.shape[1]
+    status_msg = ""
+    if done != total:
+        status_msg = f"{done}/{total} ({done / total * 100:.1f}%) indexed"
+
+    with gr.Blocks() as app:
+        
+        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(
+                        label="Taxonomic Rank",
+                        info="Which taxonomic rank to predict. Fine-grained ranks (genus, species) are more challenging.",
+                        choices=ranks,
+                        value="Species",
+                        type="index",
+                    )
+                    open_domain_btn = gr.Button("Submit", variant="primary")
+                with gr.Column():
+                    open_domain_output = gr.Label(
+                        num_top_classes=k,
+                        label="Prediction",
+                        show_label=True,
+                        value=None,
+                    )
+
+            with gr.Row():
+                gr.Examples(
+                    examples=open_domain_examples,
+                    inputs=[img_input, rank_dropdown],
+                    cache_examples=True,
+                    fn=lambda img, rank: open_domain_classification(img, rank, return_all=False),
+                    outputs=[open_domain_output],
+                )
+
+        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(
+                        placeholder="Canis familiaris (dog)\nFelis catus (cat)\n...",
+                        lines=3,
+                        label="Classes",
+                        show_label=True,
+                        info="Use taxonomic names where possible; include common names if possible.",
+                    )
+                    zero_shot_btn = gr.Button("Submit", variant="primary")
+
+                with gr.Column():
+                    zero_shot_output = gr.Label(
+                        num_top_classes=k, label="Prediction", show_label=True
+                    )
+
+            with gr.Row():
+                gr.Examples(
+                    examples=zero_shot_examples,
+                    inputs=[img_input_zs, classes_txt],
+                    cache_examples=True,
+                    fn=zero_shot_classification,
+                    outputs=[zero_shot_output],
+                )
+
+        rank_dropdown.change(
+            fn=change_output, inputs=rank_dropdown, outputs=[open_domain_output]
+        )
+
+        open_domain_btn.click(
+            fn=lambda img, rank: open_domain_classification(img, rank, return_all=True),
+            inputs=[img_input, rank_dropdown],
+            outputs=[open_domain_output, sample_img, taxon_url],
+        )
+
+        zero_shot_btn.click(
+            fn=zero_shot_classification,
+            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 2 Model](https://huggingface.co/imageomics/bioclip-2) creation, see our [BioCLIP 2 Project website](https://imageomics.github.io/bioclip-2/),
+            and for easier integration of BioCLIP 2, checkout [pybioclip](https://github.com/Imageomics/pybioclip).
+            
+            To learn more about the data, check out our [TreeOfLife-200M Dataset](https://huggingface.co/datasets/imageomics/TreeOfLife-200M).
+            """
+        )
+
+    app.queue(max_size=20)
+    app.launch(share=True)

components/metadata.parquet

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

components/query.py

@@ -0,0 +1,133 @@
+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/"
+GBIF_URL = "https://gbif.org/species/"
+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 GBIF or 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.
+    ref_page : str
+        URL to GBIF or 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, gbif_id, 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 GBIF and 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-200m-sample-images',
+                                                           'Key': filepath}
+                                                   )
+        img_resp = requests.get(img_src)
+        img = Image.open(io.BytesIO(img_resp.content))
+        if gbif_id:
+            gbif_url = GBIF_URL + gbif_id
+            if eol_page_id:
+                eol_url = EOL_URL + eol_page_id
+                if is_exact:
+                    ref_page = f"<p>Check out the <a href={eol_url} target='_blank'>EOL</a> or <a href={gbif_url} target='_blank'>GBIF</a> entry for {pred_taxon} to learn more.</p>"
+                else:
+                    ref_page = f"<p>Check out an example entry within {pred_taxon} to learn more: {full_name} at <a href={eol_url} target='_blank'>EOL</a> or <a href={gbif_url} target='_blank'>GBIF</a>.</p>"
+            else:
+                if is_exact:
+                    ref_page = f"<p>Check out the <a href={gbif_url} target='_blank'>GBIF</a> entry for {pred_taxon} to learn more.</p>"
+                else:
+                    ref_page = f"<p>Check out an example GBIF entry within {pred_taxon} to learn more: <a href={gbif_url} target='_blank'>{full_name}</a>.</p>"
+        else:
+            eol_url = EOL_URL + eol_page_id
+            if is_exact:
+                    ref_page = f"<p>Check out the <a href={eol_url} target='_blank'>EOL</a> entry for {pred_taxon} to learn more.</p>"
+            else:
+                ref_page = f"<p>Check out an example EOL entry within {pred_taxon} to learn more: <a href={eol_url} target='_blank'>{full_name}</a>.</p>"
+        logger.info(f"Successfully retrieved sample image and page for {pred_taxon}")
+        return img, ref_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.
+    gbif_id: str
+        GBIF page ID associated with predicted taxon for more information.
+    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.copy()
+    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["gbif_id"].cast(pl.String)[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["gbif_id"].cast(pl.String)[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-200m-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

@@ -0,0 +1,8 @@
+open_clip_torch
+torchvision
+torch
+gradio
+polars
+pillow
+boto3
+pybioclip