File size: 12,988 Bytes
6e5adf0
 
2cfb891
290c238
293f004
290c238
d1c1a86
2cfb891
e592678
d1c1a86
 
5cfebb1
d1c1a86
 
8fa75cc
e592678
d1c1a86
293f004
 
 
 
290c238
 
e592678
 
 
 
 
 
216fbaf
 
6e5adf0
 
 
 
 
 
216fbaf
290c238
d1c1a86
 
 
 
2cfb891
d1c1a86
 
 
 
 
 
 
2cfb891
 
6e5adf0
 
 
 
 
 
2cfb891
d4005aa
 
 
 
 
 
2cfb891
 
 
 
 
 
 
 
 
 
 
 
d4005aa
 
 
 
 
 
 
 
 
 
 
 
2cfb891
 
d1c1a86
6e5adf0
 
 
 
d1c1a86
 
 
 
 
290c238
d1c1a86
 
 
 
 
 
 
 
 
2cfb891
 
d1c1a86
 
290c238
d1c1a86
 
290c238
5cfebb1
290c238
d1c1a86
 
 
ef07580
 
 
 
 
 
 
2cfb891
e592678
290c238
6e5adf0
 
 
290c238
e592678
 
290c238
 
 
 
6e5adf0
 
2cfb891
6e5adf0
 
e592678
ef07580
6e5adf0
e592678
 
 
 
 
 
 
2cfb891
6e5adf0
 
ef07580
2cfb891
6e5adf0
e592678
 
 
 
 
 
 
 
2cfb891
e592678
 
 
290c238
 
6e5adf0
 
290c238
 
d1c1a86
293f004
216fbaf
290c238
293f004
d1c1a86
 
293f004
d1c1a86
216fbaf
d1c1a86
6e5adf0
 
 
2cfb891
 
 
 
 
 
 
 
e592678
6e5adf0
e592678
 
 
 
 
 
 
 
 
 
 
 
 
 
6e5adf0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e592678
6e5adf0
 
 
 
 
 
e592678
6e5adf0
d4005aa
e592678
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6e5adf0
e592678
 
 
6e5adf0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e592678
6e5adf0
 
 
 
e592678
6e5adf0
 
 
 
e592678
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6e5adf0
 
 
d4005aa
6e5adf0
e592678
6e5adf0
e592678
6e5adf0
d4005aa
2cfb891
 
e592678
2cfb891
 
e592678
 
 
 
 
 
 
 
 
 
2cfb891
 
e592678
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
import collections
import heapq
import json
import os
import logging

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

from templates import openai_imagenet_template
from components.query import  get_sample

log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
logging.basicConfig(level=logging.INFO, format=log_format)
logger = logging.getLogger()

hf_token = os.getenv("HF_TOKEN")

# For sample images
METADATA_PATH = "components/metadata.csv"
# Read page ID as int and filter out smaller ablation duplicated training split
metadata_df = pl.read_csv(METADATA_PATH, low_memory = False)
metadata_df = metadata_df.with_columns(pl.col("eol_page_id").cast(pl.Int64))

model_str = "hf-hub:imageomics/bioclip"
tokenizer_str = "ViT-B-16"

txt_emb_npy = "txt_emb_species.npy"
txt_names_json = "txt_emb_species.json"

min_prob = 1e-9
k = 5

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


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()]
    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):
    taxon = " ".join(taxon)
    if not common:
        return taxon
    return f"{taxon} ({common})"


@torch.no_grad()
def open_domain_classification(img, rank: int, return_all=False):
    """
    Predicts from the entire tree of life.
    If targeting a higher rank than species, then this function predicts among all
    species, then sums up species-level probabilities for the given rank.
    """

    logger.info(f"Starting open domain classification for rank: {rank}")
    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_emb).squeeze()
    probs = F.softmax(logits, dim=0)

    if rank + 1 == len(ranks):
        topk = probs.topk(k)
        prediction_dict = {
            format_name(*txt_names[i]): prob for i, prob in zip(topk.indices, topk.values)
        }
        logger.info(f"Top K predictions: {prediction_dict}")
        top_prediction_name = format_name(*txt_names[topk.indices[0]]).split("(")[0]
        logger.info(f"Top prediction name: {top_prediction_name}")
        sample_img, taxon_url = get_sample(metadata_df, top_prediction_name, rank)
        if return_all:
            return prediction_dict, sample_img, taxon_url
        return prediction_dict

    output = collections.defaultdict(float)
    for i in torch.nonzero(probs > min_prob).squeeze():
        output[" ".join(txt_names[i][0][: rank + 1])] += probs[i]

    topk_names = heapq.nlargest(k, output, key=output.get)
    prediction_dict = {name: output[name] for name in topk_names}
    logger.info(f"Top K names for output: {topk_names}")
    logger.info(f"Prediction dictionary: {prediction_dict}")

    top_prediction_name = topk_names[0]
    logger.info(f"Top prediction name: {top_prediction_name}")
    sample_img, taxon_url = get_sample(metadata_df, top_prediction_name, rank)
    logger.info(f"Sample image and taxon URL: {sample_img}, {taxon_url}")

    if return_all:
        return prediction_dict, sample_img, taxon_url
    return prediction_dict


def change_output(choice):
    return gr.Label(num_top_classes=k, label=ranks[choice], show_label=True, value=None)


if __name__ == "__main__":
    logger.info("Starting.")
    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)

    txt_emb = torch.from_numpy(np.load(txt_emb_npy, mmap_mode="r")).to(device)
    with open(txt_names_json) as fd:
        txt_names = json.load(fd)

    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:
        
        with gr.Tab("Open-Ended"):
            with gr.Row(variant = "panel", elem_id = "images_panel"):
                with gr.Column():
                    img_input = gr.Image(height = 400, sources=["upload"])
            
                with gr.Column():
                    # display sample image of top predicted taxon
                    sample_img = gr.Image(label = "Sample Image of Predicted Taxon", 
                                        height = 400, 
                                        show_download_button = False)
                
                    taxon_url = gr.HTML(label = "More Information", 
                                    elem_id = "url"
                                    )

            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")
                with gr.Column():
                    open_domain_output = gr.Label(
                        num_top_classes=k,
                        label="Prediction",
                        show_label=True,
                        value=None,
                    )
                  #  open_domain_flag_btn = gr.Button("Flag Mistake", variant="primary")

            with gr.Row():
                gr.Examples(
                    examples=open_domain_examples,
                    inputs=[img_input, rank_dropdown],
                    cache_examples=True,
                    fn=lambda img, rank: open_domain_classification(img, rank, return_all=False),
                    outputs=[open_domain_output],
                )
            '''
            # Flagging Code
            open_domain_callback = gr.HuggingFaceDatasetSaver(
                hf_token, "bioclip-demo-open-domain-mistakes", private=True
            )
            open_domain_callback.setup(
                [img_input, rank_dropdown, open_domain_output],
                flagging_dir="bioclip-demo-open-domain-mistakes/logs/flagged",
            )
            open_domain_flag_btn.click(
                lambda *args: open_domain_callback.flag(args),
                [img_input, rank_dropdown, open_domain_output],
                None,
                preprocess=False,
            )
            '''
        with gr.Tab("Zero-Shot"):
            with gr.Row():
                img_input_zs = gr.Image(height = 400, sources=["upload"])

            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=k, 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_zs, classes_txt],
                    cache_examples=True,
                    fn=zero_shot_classification,
                    outputs=[zero_shot_output],
                )
        '''
        # Flagging Code
        zero_shot_callback = gr.HuggingFaceDatasetSaver(
            hf_token, "bioclip-demo-zero-shot-mistakes", private=True
        )
        zero_shot_callback.setup(
            [img_input, zero_shot_output], flagging_dir="bioclip-demo-zero-shot-mistakes/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_output]
        )

        open_domain_btn.click(
            fn=lambda img, rank: open_domain_classification(img, rank, return_all=True),
            inputs=[img_input, rank_dropdown],
            outputs=[open_domain_output, sample_img, taxon_url],
        )

        zero_shot_btn.click(
            fn=zero_shot_classification,
            inputs=[img_input_zs, classes_txt],
            outputs=zero_shot_output,
        )
        
        # Footer to point out to model and data from app page.
        gr.Markdown(
            """
            For more information on the [BioCLIP Model](https://huggingface.co/imageomics/bioclip) creation, see our [BioCLIP Project GitHub](https://github.com/Imageomics/bioclip), and
            for easier integration of BioCLIP, checkout [pybioclip](https://github.com/Imageomics/pybioclip).
            
            To learn more about the data, check out our [TreeOfLife-10M Dataset](https://huggingface.co/datasets/imageomics/TreeOfLife-10M).
            """
        )

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