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