revert back to local inference code

Co-authored-by: Jianyang Gu <[email protected]>

app.py

@@ -11,9 +11,10 @@ import torch
 import torch.nn.functional as F
 from open_clip import create_model, get_tokenizer
 from torchvision import transforms
+from huggingface_hub import hf_hub_download
 
+from components.templates import openai_imagenet_template
 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)
@@ -27,16 +28,16 @@ METADATA_PATH = "components/metadata.parquet"
 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"
+model_str = "hf-hub:imageomics/bioclip-2"
+tokenizer_str = "ViT-L-14"
+HF_DATA_STR = "imageomics/TreeOfLife-200M"
 
-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"
+txt_names_json = "components/txt_emb_species.json"
 
 min_prob = 1e-9
 k = 5
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+device = torch.device("cpu")
 
 preprocess_img = transforms.Compose(
     [
@@ -52,41 +53,45 @@ 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"],
-    ["examples/Sarcoscypha-coccinea.jpeg", "Order"],
+    ["examples/Carcharhinus-melanopterus.jpg", "Species"],
+    ["examples/house-finch.jpeg", "Species"],
+    ["examples/Bovidae-Oryx.jpg", "Genus"],
+    ["examples/Cebidae-Cebus.jpg", "Genus"],
+    ["examples/Solanales-Petunia.png", "Genus"],
+    ["examples/Asparagales-Orchidaceae.jpg", "Family"],
 ]
 zero_shot_examples = [
     [
-        "examples/Ursus-arctos.jpeg",
-        "brown bear\nblack bear\npolar bear\nkoala bear\ngrizzly bear",
+        "examples/Cortinarius-austroalbidus.jpg",
+        "Cortinarius austroalbidus\nCortinarius armillatus\nCortinarius caperatus"
     ],
-    ["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/leopard.jpg",
+        "Jaguar\nLeopard\nCheetah",
+    ],
+    [
+        "examples/jaguar.jpg",
+        "Jaguar\nLeopard\nCheetah",
     ],
     [
-        "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/cheetah.jpg",
+        "Jaguar\nLeopard\nCheetah",
     ],
     [
-        "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/monarch.jpg",
+        "Danaus plexippus — Monarch\nLimenitis archippus — Viceroy",
     ],
     [
-        "examples/Sarcoscypha-coccinea.jpeg",
-        "scarlet elf cup (coccinea)\nscharlachroter kelchbecherling (austriaca)\ncrimson cup (dudleyi)\nstalked scarlet cup (occidentalis)",
+        "examples/viceroy.jpg",
+        "Danaus plexippus — Monarch\nLimenitis archippus — Viceroy",
     ],
     [
-        "examples/Onoclea-hintonii.jpg",
-        "Onoclea attenuata\nOnoclea boryana\nOnoclea hintonii\nOnoclea intermedia\nOnoclea sensibilis",
+        "examples/Ursus-arctos.jpeg",
+        "brown bear\nblack bear\npolar bear\nkoala bear\ngrizzly bear",
     ],
     [
-        "examples/Onoclea-sensibilis.jpg",
-        "Onoclea attenuata\nOnoclea boryana\nOnoclea hintonii\nOnoclea intermedia\nOnoclea sensibilis",
+        "examples/Carnegiea-gigantea.png",
+        "Carnegiea gigantea\nSchlumbergera opuntioides\nMammillaria albicoma",
     ],
 ]
 
@@ -95,13 +100,32 @@ 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()]
-    classifier = CustomLabelsClassifier(
-        cls_ary = classes,
-        model_str = MODEL_STR,  # remove this line once pybioclip uses BioCLIP 2
-    )
-    return classifier.predict(img)
+    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):
@@ -165,16 +189,20 @@ def change_output(choice):
 
 if __name__ == "__main__":
     logger.info("Starting.")
-    model = create_model(MODEL_STR, output_dict=True, require_pretrained=True)
+    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)
+    tokenizer = get_tokenizer(tokenizer_str)
 
-    txt_emb = torch.from_numpy(np.load(txt_emb_npy, mmap_mode="r")).to(device)
+    txt_emb = torch.from_numpy(np.load(hf_hub_download(
+        repo_id=HF_DATA_STR,
+        filename="embeddings/txt_emb_species.npy",
+        repo_type="dataset",
+    )))
     with open(txt_names_json) as fd:
         txt_names = json.load(fd)
 

components/templates.py

@@ -0,0 +1,82 @@
+openai_imagenet_template = [
+    lambda c: f"a bad photo of a {c}.",
+    lambda c: f"a photo of many {c}.",
+    lambda c: f"a sculpture of a {c}.",
+    lambda c: f"a photo of the hard to see {c}.",
+    lambda c: f"a low resolution photo of the {c}.",
+    lambda c: f"a rendering of a {c}.",
+    lambda c: f"graffiti of a {c}.",
+    lambda c: f"a bad photo of the {c}.",
+    lambda c: f"a cropped photo of the {c}.",
+    lambda c: f"a tattoo of a {c}.",
+    lambda c: f"the embroidered {c}.",
+    lambda c: f"a photo of a hard to see {c}.",
+    lambda c: f"a bright photo of a {c}.",
+    lambda c: f"a photo of a clean {c}.",
+    lambda c: f"a photo of a dirty {c}.",
+    lambda c: f"a dark photo of the {c}.",
+    lambda c: f"a drawing of a {c}.",
+    lambda c: f"a photo of my {c}.",
+    lambda c: f"the plastic {c}.",
+    lambda c: f"a photo of the cool {c}.",
+    lambda c: f"a close-up photo of a {c}.",
+    lambda c: f"a black and white photo of the {c}.",
+    lambda c: f"a painting of the {c}.",
+    lambda c: f"a painting of a {c}.",
+    lambda c: f"a pixelated photo of the {c}.",
+    lambda c: f"a sculpture of the {c}.",
+    lambda c: f"a bright photo of the {c}.",
+    lambda c: f"a cropped photo of a {c}.",
+    lambda c: f"a plastic {c}.",
+    lambda c: f"a photo of the dirty {c}.",
+    lambda c: f"a jpeg corrupted photo of a {c}.",
+    lambda c: f"a blurry photo of the {c}.",
+    lambda c: f"a photo of the {c}.",
+    lambda c: f"a good photo of the {c}.",
+    lambda c: f"a rendering of the {c}.",
+    lambda c: f"a {c} in a video game.",
+    lambda c: f"a photo of one {c}.",
+    lambda c: f"a doodle of a {c}.",
+    lambda c: f"a close-up photo of the {c}.",
+    lambda c: f"a photo of a {c}.",
+    lambda c: f"the origami {c}.",
+    lambda c: f"the {c} in a video game.",
+    lambda c: f"a sketch of a {c}.",
+    lambda c: f"a doodle of the {c}.",
+    lambda c: f"a origami {c}.",
+    lambda c: f"a low resolution photo of a {c}.",
+    lambda c: f"the toy {c}.",
+    lambda c: f"a rendition of the {c}.",
+    lambda c: f"a photo of the clean {c}.",
+    lambda c: f"a photo of a large {c}.",
+    lambda c: f"a rendition of a {c}.",
+    lambda c: f"a photo of a nice {c}.",
+    lambda c: f"a photo of a weird {c}.",
+    lambda c: f"a blurry photo of a {c}.",
+    lambda c: f"a cartoon {c}.",
+    lambda c: f"art of a {c}.",
+    lambda c: f"a sketch of the {c}.",
+    lambda c: f"a embroidered {c}.",
+    lambda c: f"a pixelated photo of a {c}.",
+    lambda c: f"itap of the {c}.",
+    lambda c: f"a jpeg corrupted photo of the {c}.",
+    lambda c: f"a good photo of a {c}.",
+    lambda c: f"a plushie {c}.",
+    lambda c: f"a photo of the nice {c}.",
+    lambda c: f"a photo of the small {c}.",
+    lambda c: f"a photo of the weird {c}.",
+    lambda c: f"the cartoon {c}.",
+    lambda c: f"art of the {c}.",
+    lambda c: f"a drawing of the {c}.",
+    lambda c: f"a photo of the large {c}.",
+    lambda c: f"a black and white photo of a {c}.",
+    lambda c: f"the plushie {c}.",
+    lambda c: f"a dark photo of a {c}.",
+    lambda c: f"itap of a {c}.",
+    lambda c: f"graffiti of the {c}.",
+    lambda c: f"a toy {c}.",
+    lambda c: f"itap of my {c}.",
+    lambda c: f"a photo of a cool {c}.",
+    lambda c: f"a photo of a small {c}.",
+    lambda c: f"a tattoo of the {c}.",
+]

components/txt_emb_species.json

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