Spaces:
Running
Running
File size: 10,127 Bytes
6e5adf0 2cfb891 290c238 293f004 290c238 d1c1a86 2cfb891 d1c1a86 5cfebb1 d1c1a86 8fa75cc d1c1a86 293f004 290c238 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 6e5adf0 290c238 6e5adf0 290c238 6e5adf0 2cfb891 6e5adf0 ef07580 6e5adf0 2cfb891 6e5adf0 ef07580 2cfb891 6e5adf0 2cfb891 6e5adf0 290c238 6e5adf0 290c238 d1c1a86 293f004 216fbaf 290c238 293f004 d1c1a86 293f004 d1c1a86 216fbaf d1c1a86 6e5adf0 2cfb891 293f004 d4005aa 6e5adf0 d4005aa 6e5adf0 2cfb891 6e5adf0 d4005aa 2cfb891 d4005aa 2cfb891 6e5adf0 d4005aa 6e5adf0 d4005aa 2cfb891 |
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 |
import collections
import heapq
import json
import os
import logging
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
from templates import openai_imagenet_template
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")
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) -> dict[str, float]:
"""
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.
"""
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 predicting species, no need to sum probabilities.
if rank + 1 == len(ranks):
topk = probs.topk(k)
return {
format_name(*txt_names[i]): prob for i, prob in zip(topk.indices, topk.values)
}
# Sum up by the rank
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)
return {name: output[name] for name in topk_names}
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:
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")
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=open_domain_classification,
outputs=[open_domain_output],
)
open_domain_callback = gr.HuggingFaceDatasetSaver(
hf_token, "imageomics/bioclip-demo-open-domain-mistakes", private=True
)
open_domain_callback.setup(
[img_input, rank_dropdown, open_domain_output],
flagging_dir="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():
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, classes_txt],
cache_examples=True,
fn=zero_shot_classification,
outputs=[zero_shot_output],
)
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_output]
)
open_domain_btn.click(
fn=open_domain_classification,
inputs=[img_input, rank_dropdown],
outputs=[open_domain_output],
)
zero_shot_btn.click(
fn=zero_shot_classification,
inputs=[img_input, classes_txt],
outputs=zero_shot_output,
)
app.queue(max_size=20)
app.launch()
|