Add sample image return functionality to app

pulls sample image from EOL subset of TreeOfLife-10M

app.py

@@ -6,12 +6,14 @@ 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)
@@ -19,6 +21,12 @@ 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"
 
@@ -123,12 +131,14 @@ def format_name(taxon, common):
 
 
 @torch.no_grad()
-def open_domain_classification(img, rank: int) -> dict[str, float]:
+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)
@@ -136,21 +146,36 @@ def open_domain_classification(img, rank: int) -> dict[str, float]:
  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 {
+ 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
 
- # 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)
+ 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}")
 
- return {name: output[name] for name in topk_names}
+ if return_all:
+ return prediction_dict, sample_img, taxon_url
+ return prediction_dict
 
 
 def change_output(choice):
@@ -179,9 +204,22 @@ 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(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(
@@ -199,32 +237,36 @@ if __name__ == "__main__":
  show_label=True,
  value=None,
  )
- # open_domain_flag_btn = gr.Button("Flag Mistake", 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=open_domain_classification,
+ fn=lambda img, rank: open_domain_classification(img, rank, return_all=False),
  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,
- # )
-
+ '''
+ # 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(
@@ -240,43 +282,44 @@ if __name__ == "__main__":
  zero_shot_output = gr.Label(
  num_top_classes=k, label="Prediction", show_label=True
  )
- # zero_shot_flag_btn = gr.Button("Flag Mistake", variant="primary")
+ # zero_shot_flag_btn = gr.Button("Flag Mistake", variant="primary")
 
  with gr.Row():
  gr.Examples(
  examples=zero_shot_examples,
- inputs=[img_input, classes_txt],
+ inputs=[img_input_zs, 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,
- # )
-
+ '''
+ # 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=open_domain_classification,
+ fn=lambda img, rank: open_domain_classification(img, rank, return_all=True),
  inputs=[img_input, rank_dropdown],
- outputs=[open_domain_output],
+ outputs=[open_domain_output, sample_img, taxon_url],
  )
 
  zero_shot_btn.click(
  fn=zero_shot_classification,
- inputs=[img_input, classes_txt],
+ inputs=[img_input_zs, classes_txt],
  outputs=zero_shot_output,
  )
  
@@ -291,4 +334,4 @@ if __name__ == "__main__":
  )
 
  app.queue(max_size=20)
- app.launch()
+ app.launch(share=True)