v0.1

README.md

@@ -1,11 +1,11 @@
 ---
 title: Bioclip Demo
-emoji: 👀
+emoji: 🐘
 colorFrom: indigo
 colorTo: purple
 sdk: gradio
 sdk_version: 4.7.1
 app_file: app.py
-pinned: false
+pinned: true
 license: mit
 ---

app.py

@@ -1,24 +1,29 @@
+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")
-hf_writer = gr.HuggingFaceDatasetSaver(hf_token, "bioclip-demo")
 
 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),
@@ -26,6 +31,28 @@ preprocess_img = transforms.Compose(
     ]
 )
 
+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):
@@ -42,8 +69,8 @@ def get_txt_features(classnames, templates):
 
 
 @torch.no_grad()
-def predict(img, classes: list[str]) -> dict[str, float]:
-    classes = [cls.strip() for cls in classes if cls.strip()]
+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)
@@ -55,7 +82,8 @@ def predict(img, classes: list[str]) -> dict[str, float]:
     return {cls: prob for cls, prob in zip(classes, probs)}
 
 
-def hierarchical_predict(img) -> list[str]:
+@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.
     """
@@ -63,16 +91,44 @@ def hierarchical_predict(img) -> list[str]:
     img_features = model.encode_image(img.unsqueeze(0))
     img_features = F.normalize(img_features, dim=-1)
 
-    breakpoint()
+    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 run(img, cls_str: str) -> dict[str, float]:
-    breakpoint()
-    if cls_str:
-        classes = [cls.strip() for cls in cls_str.split("\n") if cls.strip()]
-        return predict(img, classes)
-    else:
-        return hierarchical_predict(img)
+def get_name_lookup(path):
+    with open(path) as fd:
+        return lib.TaxonomicTree.from_dict(json.load(fd))
 
 
 if __name__ == "__main__":
@@ -86,22 +142,104 @@ if __name__ == "__main__":
 
     tokenizer = get_tokenizer(tokenizer_str)
 
-    demo = gr.Interface(
-        fn=run,
-        inputs=[
-            gr.Image(shape=(224, 224)),
-            gr.Textbox(
-                placeholder="dog\ncat\n...",
-                lines=3,
-                label="Classes",
-                show_label=True,
-                info="If empty, will predict from the entire tree of life.",
-            ),
-        ],
-        outputs=gr.Label(num_top_classes=20, label="Predictions", show_label=True),
-        allow_flagging="manual",
-        flagging_options=["Incorrect", "Other"],
-        flagging_callback=hf_writer,
-    )
-
-    demo.launch()
+    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()

lib.py

@@ -1,5 +1,5 @@
-import json
 import itertools
+import json
 
 
 class TaxonomicNode:
@@ -43,11 +43,12 @@ class TaxonomicNode:
     @classmethod
     def from_dict(cls, dct, root):
         node = cls(dct["name"], dct["index"], root)
-        node._children = {child["name"]: cls.from_dict(child, root) for child in dct["children"]}
+        node._children = {
+            child["name"]: cls.from_dict(child, root) for child in dct["children"]
+        }
         return node
 
 
-
 class TaxonomicTree:
     """
     Efficient structure for finding taxonomic names and their descendants.
@@ -85,11 +86,15 @@ class TaxonomicTree:
         for kingdom in self.kingdoms.values():
             yield from kingdom
 
+    def __len__(self):
+        return self.size
+
     @classmethod
     def from_dict(cls, dct):
         tree = cls()
         tree.kingdoms = {
-            kingdom["name"]: TaxonomicNode.from_dict(kingdom, tree) for kingdom in dct["kingdoms"]
+            kingdom["name"]: TaxonomicNode.from_dict(kingdom, tree)
+            for kingdom in dct["kingdoms"]
         }
         tree.size = dct["size"]
         return tree
@@ -112,11 +117,10 @@ class TaxonomicJsonEncoder(json.JSONEncoder):
             super().default(self, obj)
 
 
-
 def batched(iterable, n):
     # batched('ABCDEFG', 3) --> ABC DEF G
     if n < 1:
-        raise ValueError('n must be at least one')
+        raise ValueError("n must be at least one")
     it = iter(iterable)
     while batch := tuple(itertools.islice(it, n)):
-        yield zip(*batch)
+        yield zip(*batch)

make_txt_embedding.py

@@ -5,6 +5,7 @@ Uses the catalog.csv file from TreeOfLife-10M.
 import argparse
 import csv
 import json
+import os
 
 import numpy as np
 import torch
@@ -22,29 +23,53 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 @torch.no_grad()
 def write_txt_features(name_lookup):
-    all_features = np.memmap(
-        args.out_path, dtype=np.float32, mode="w+", shape=(512, name_lookup.size)
-    )
+    if os.path.isfile(args.out_path):
+        all_features = np.load(args.out_path)
+    else:
+        all_features = np.zeros((512, len(name_lookup)), dtype=np.float32)
 
     batch_size = args.batch_size // len(openai_imagenet_template)
-    for names, indices in tqdm(lib.batched(name_lookup, batch_size)):
-        txts = [template(name) for name in names for template in openai_imagenet_template]
+    for batch, (names, indices) in enumerate(
+        tqdm(
+            lib.batched(name_lookup, batch_size),
+            desc="txt feats",
+            total=len(name_lookup) // batch_size,
+        )
+    ):
+        # Skip if any non-zero elements
+        if all_features[:, indices].any():
+            print(f"Skipping batch {batch}")
+            continue
+
+        txts = [
+            template(name) for name in names for template in openai_imagenet_template
+        ]
         txts = tokenizer(txts).to(device)
         txt_features = model.encode_text(txts)
-        txt_features = torch.reshape(txt_features, (batch_size, len(openai_imagenet_template), 512))
+        txt_features = torch.reshape(
+            txt_features, (len(names), len(openai_imagenet_template), 512)
+        )
         txt_features = F.normalize(txt_features, dim=2).mean(dim=1)
         txt_features /= txt_features.norm(dim=1, keepdim=True)
-        all_features[:, indices] = txt_features.cpu().numpy().T
+        all_features[:, indices] = txt_features.T.cpu().numpy()
+
+        if batch % 100 == 0:
+            np.save(args.out_path, all_features)
 
-    all_features.flush()
+    np.save(args.out_path, all_features)
 
 
-def get_name_lookup(catalog_path):
+def get_name_lookup(catalog_path, cache_path):
+    if os.path.isfile(cache_path):
+        with open(cache_path) as fd:
+            lookup = lib.TaxonomicTree.from_dict(json.load(fd))
+        return lookup
+
     lookup = lib.TaxonomicTree()
 
     with open(catalog_path) as fd:
         reader = csv.DictReader(fd)
-        for row in tqdm(reader):
+        for row in tqdm(reader, desc="catalog"):
             name = [
                 row["kingdom"],
                 row["phylum"],
@@ -58,6 +83,9 @@ def get_name_lookup(catalog_path):
                 name = name[: name.index("")]
             lookup.add(name)
 
+    with open(args.name_cache_path, "w") as fd:
+        json.dump(lookup, fd, cls=lib.TaxonomicJsonEncoder)
+
     return lookup
 
 
@@ -69,15 +97,17 @@ if __name__ == "__main__":
         required=True,
     )
     parser.add_argument("--out-path", help="Path to the output file.", required=True)
-    parser.add_argument("--name-cache-path", help="Path to the name cache file.", default=".name_lookup_cache.json")
-    parser.add_argument("--batch-size", help="Batch size.", default=2 ** 15, type=int)
+    parser.add_argument(
+        "--name-cache-path",
+        help="Path to the name cache file.",
+        default="name_lookup.json",
+    )
+    parser.add_argument("--batch-size", help="Batch size.", default=2**15, type=int)
     args = parser.parse_args()
 
-    name_lookup = get_name_lookup(args.catalog_path)
-    with open(args.name_cache_path, "w") as fd:
-        json.dump(name_lookup, fd, cls=lib.TaxonomicJsonEncoder)
+    name_lookup = get_name_lookup(args.catalog_path, cache_path=args.name_cache_path)
+    print("Got name lookup.")
 
-    print("Starting.")
     model = create_model(model_str, output_dict=True, require_pretrained=True)
     model = model.to(device)
     print("Created model.")

txt_emb.npy

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