more image examples; no open-ended prediction

app.py

@@ -33,13 +33,19 @@ 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"],
-]
+# 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",
@@ -52,6 +58,18 @@ zero_shot_examples = [
         "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",
+    ],
 ]
 
 
@@ -153,89 +171,94 @@ if __name__ == "__main__":
         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,
+        img_input = gr.Image(height=512)
+
+        # 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")
+
+            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")
+
+        with gr.Row():
+            gr.Examples(
+                examples=zero_shot_examples,
+                inputs=[img_input, classes_txt],
+                cache_examples=True,
+                fn=zero_shot_classification,
+                outputs=[zero_shot_output],
             )
 
-        rank_dropdown.change(
-            fn=change_output, inputs=rank_dropdown, outputs=open_domain_outputs
+        zero_shot_callback = gr.HuggingFaceDatasetSaver(
+            hf_token, "imageomics/bioclip-demo-zero-shot-mistakes", private=True
         )
-
-        open_domain_btn.click(
-            fn=open_domain_classification,
-            inputs=[img_input, rank_dropdown],
-            outputs=open_domain_outputs,
+        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],

make_txt_embedding.py

@@ -74,24 +74,31 @@ def convert_txt_features_to_avgs(name_lookup):
     all_features = torch.from_numpy(np.load(args.out_path)).to(device)
     logger.info("Loaded text features from disk to %s.", device)
 
-    all_names = [set() for rank in ranks]
+    names_by_rank = [set() for rank in ranks]
     for name, index in tqdm(name_lookup.values()):
         i = len(name) - 1
-        all_names[i].add((name, index))
+        names_by_rank[i].add((name, index))
 
     zeroed = 0
     for i, rank in reversed(list(enumerate(ranks))):
         if rank == "Species":
             continue
-        for name, index in tqdm(all_names[i], desc=rank):
-            species = tuple(zip(*((d, i) for d, i in name_lookup.descendants(prefix=name) if len(d) >= 6)))
+        for name, index in tqdm(names_by_rank[i], desc=rank):
+            species = tuple(
+                zip(
+                    *(
+                        (d, i)
+                        for d, i in name_lookup.descendants(prefix=name)
+                        if len(d) >= 6
+                    )
+                )
+            )
             if not species:
                 logger.warning("No species for %s.", " ".join(name))
                 all_features[:, index] = 0.0
                 zeroed += 1
                 continue
 
-
             values, indices = species
             mean = all_features[:, indices].mean(dim=1)
             all_features[:, index] = F.normalize(mean, dim=0)
@@ -99,7 +106,10 @@ def convert_txt_features_to_avgs(name_lookup):
     out_path, ext = os.path.splitext(args.out_path)
     np.save(f"{out_path}_avgs{ext}", all_features.cpu().numpy())
     if zeroed:
-        logger.warning("Zeroed out %d nodes because they didn't have any genus or species-level labels.", zeroed)
+        logger.warning(
+            "Zeroed out %d nodes because they didn't have any genus or species-level labels.",
+            zeroed,
+        )
 
 
 def get_name_lookup(catalog_path, cache_path):