app.py

import json
import os

import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
from open_clip import create_model, get_tokenizer
from torchvision import transforms

import lib
from templates import openai_imagenet_template

hf_token = os.getenv("HF_TOKEN")

model_str = "hf-hub:imageomics/bioclip"
tokenizer_str = "ViT-B-16"
name_lookup_json = "name_lookup.json"
txt_emb_npy = "txt_emb.npy"

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"],
]
zero_shot_examples = [
    [
        "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",
    ],
]


@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)}


@torch.no_grad()
def open_domain_classification(img, rank: int) -> list[dict[str, float]]:
    """
    Predicts from the top of the tree of life down to the species.
    """
    img = preprocess_img(img).to(device)
    img_features = model.encode_image(img.unsqueeze(0))
    img_features = F.normalize(img_features, dim=-1)

    outputs = []

    name = []
    for _ in range(rank + 1):
        children = tuple(zip(*name_lookup.children(name)))
        if not children:
            break
        values, indices = children
        txt_features = txt_emb[:, indices].to(device)
        logits = (model.logit_scale.exp() * img_features @ txt_features).view(-1)

        probs = F.softmax(logits, dim=0).to("cpu").tolist()
        parent = " ".join(name)
        outputs.append(
            {f"{parent} {value}": prob for value, prob in zip(values, probs)}
        )

        top = values[logits.argmax()]
        name.append(top)

    while len(outputs) < 7:
        outputs.append({})

    return list(reversed(outputs))


def change_output(choice):
    return [
        gr.Label(
            num_top_classes=5, label=rank, show_label=True, visible=(6 - i <= choice)
        )
        for i, rank in enumerate(reversed(ranks))
    ]


def get_name_lookup(path):
    with open(path) as fd:
        return lib.TaxonomicTree.from_dict(json.load(fd))


if __name__ == "__main__":
    print("Starting.")
    model = create_model(model_str, output_dict=True, require_pretrained=True)
    model = model.to(device)
    print("Created model.")

    model = torch.compile(model)
    print("Compiled model.")

    tokenizer = get_tokenizer(tokenizer_str)

    name_lookup = get_name_lookup(name_lookup_json)
    txt_emb = torch.from_numpy(np.load(txt_emb_npy, mmap_mode="r"))

    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:
        img_input = gr.Image()

        with gr.Tab("Open-Ended"):
            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")
                    gr.Examples(
                        examples=open_domain_examples,
                        inputs=[img_input, rank_dropdown],
                    )

                with gr.Column():
                    open_domain_outputs = [
                        gr.Label(num_top_classes=5, label=rank, show_label=True)
                        for rank in reversed(ranks)
                    ]
                    open_domain_flag_btn = gr.Button("Flag Mistake", variant="primary")

            open_domain_callback = gr.HuggingFaceDatasetSaver(
                hf_token, "imageomics/bioclip-demo-open-domain-mistakes", private=True
            )
            open_domain_callback.setup(
                [img_input, *open_domain_outputs], flagging_dir="logs/flagged"
            )
            open_domain_flag_btn.click(
                lambda *args: open_domain_callback.flag(args),
                [img_input, *open_domain_outputs],
                None,
                preprocess=False,
            )

        with gr.Tab("Zero-Shot"):
            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")
                    gr.Examples(
                        examples=zero_shot_examples,
                        inputs=[img_input, classes_txt],
                    )

                with gr.Column():
                    zero_shot_output = gr.Label(
                        num_top_classes=5, label="Prediction", show_label=True
                    )
                    zero_shot_flag_btn = gr.Button("Flag Mistake", variant="primary")

            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,
            )

        rank_dropdown.change(
            fn=change_output, inputs=rank_dropdown, outputs=open_domain_outputs
        )

        open_domain_btn.click(
            fn=open_domain_classification,
            inputs=[img_input, rank_dropdown],
            outputs=open_domain_outputs,
        )

        zero_shot_btn.click(
            fn=zero_shot_classification,
            inputs=[img_input, classes_txt],
            outputs=zero_shot_output,
        )

    app.queue(max_size=20)
    app.launch()