reefnet-demo / app.py
Yahia battach
edit app.py
4e412dc
raw
history blame
12.2 kB
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/Acropora-gemmifera.jpg", "Species"],
["examples/Favites_abdita.jpg", "Genus"],
["examples/Pocillopora_acuta.jpg", "Family"],
["examples/Acropora_millepora.jpg", "Species"],
["examples/porities_lobata.jpg", "Species"],
["examples/Fungia_concinna.jpg", "Family"],
]
zero_shot_examples = [
[
"examples/Acropora-gemmifera.jpg",
"Acropora aculeus\nAcropora acuminata\nAcropora anthocercis\nAcropora appressa\nAcropora arabensis"
],
["examples/porities_lobata.jpg", "Porities lobata\nPorites astreoides"],
["examples/Euphyllia_paraancora.jpg", "Euphyllia paraancora\nEuphyllia paradivisa"],
[
"examples/Montipora_patula.jpg",
"Montipora patula\nMontipora peltiformis \nMontipora saudii\nMontipora caliculata\nMontipora capitata\nMontipora cebuensis\nMontipora carinata",
],
[
"examples/Astreopora_listeri.jpg",
"Animalia Cnidaria Anthozoa Scleractinia Acroporidae Acropora hyacinthus\nAnimalia Cnidaria Anthozoa Scleractinia Acroporidae Acropora aculeus\nAnimalia Cnidaria Anthozoa Scleractinia Acroporidae Acropora anthocercis\nAnimalia Cnidaria Anthozoa Scleractinia Acroporidae Acropora millepora\nAnimalia Cnidaria Anthozoa Scleractinia Acroporidae Acropora gemmifera",
],
[
"examples/Turbinaria_heronensis.jpg",
"Turbinaria heronensis\nTurbinaria mesenterina\nTurbinaria patula\nTurbinaria peltata",
],
[
"examples/Montipora_peltiformis.jpg",
"Montipora peltiformis\nMontipora capricornis\nMontipora carinata\nMontipora cebuensis\nMontipora circumvallata",
],
[
"examples/Agaricia_agaricites.jpg",
"Agaricia agaricites\nAgaricia fragilis\nAgaricia grahamae\nAgaricia humilis\nAgaricia lamarcki",
],
]
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"
# )
open_domain_output = gr.Label(
num_top_classes=k,
label="Prediction",
show_label=True,
value=None,
)
with gr.Row():
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")
# 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="Montipora peltiformis \nMontipora saudii...",
lines=3,
label="Classes",
show_label=True,
info="Use taxonomic names",
)
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,
)
app.queue(max_size=20)
app.launch(share=True)