added app.py

README.md

@@ -1,5 +1,5 @@
 ---
-title: Bioclip Demo
+title: ReefNet Demo
 emoji: 🐘
 colorFrom: indigo
 colorTo: purple

app.py

@@ -0,0 +1,337 @@
+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 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)
+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"
+
+txt_emb_npy = "txt_emb_species.npy"
+txt_names_json = "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]
+
+
+@torch.no_grad()
+def get_txt_features(classnames, templates):
+    all_features = []
+    for classname in classnames:
+        txts = [template(classname) for template in templates]
+        txts = tokenizer(txts).to(device)
+        txt_features = model.encode_text(txts)
+        txt_features = F.normalize(txt_features, dim=-1).mean(dim=0)
+        txt_features /= txt_features.norm()
+        all_features.append(txt_features)
+    all_features = torch.stack(all_features, dim=1)
+    return all_features
+
+
+@torch.no_grad()
+def zero_shot_classification(img, cls_str: str) -> dict[str, float]:
+    classes = [cls.strip() for cls in cls_str.split("\n") if cls.strip()]
+    txt_features = get_txt_features(classes, openai_imagenet_template)
+
+    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_features).squeeze()
+    probs = F.softmax(logits, dim=0).to("cpu").tolist()
+    return {cls: prob for cls, prob in zip(classes, probs)}
+
+
+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,
+                    )
+                  #  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=lambda img, rank: open_domain_classification(img, rank, return_all=False),
+                    outputs=[open_domain_output],
+                )
+            '''
+            # 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(
+                        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
+                    )
+             #       zero_shot_flag_btn = gr.Button("Flag Mistake", variant="primary")
+
+            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],
+                )
+        '''
+        # 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=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 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(share=True)