edit app.py

app.py

@@ -129,6 +129,53 @@ def format_name(taxon, common):
  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
+
 @torch.no_grad()
 def open_domain_classification(img, rank: int, return_all=False):
  """
@@ -136,7 +183,6 @@ def open_domain_classification(img, rank: int, return_all=False):
  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))
@@ -148,15 +194,13 @@ def open_domain_classification(img, rank: int, return_all=False):
  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)
+ format_name(*txt_names[i]): prob.item() 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
+ return prediction_dict, None, None # Return dummy None values for unused parts
+ return prediction_dict # Only return the dictionary for the Label component
 
  output = collections.defaultdict(float)
  for i in torch.nonzero(probs > min_prob).squeeze():
@@ -165,18 +209,11 @@ def open_domain_classification(img, rank: int, return_all=False):
  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, None, None
  return prediction_dict
 
-
 def change_output(choice):
  return gr.Label(num_top_classes=k, label=ranks[choice], show_label=True, value=None)
 
@@ -310,12 +347,19 @@ if __name__ == "__main__":
  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],
+ # )
+
  open_domain_btn.click(
- fn=lambda img, rank: open_domain_classification(img, rank, return_all=True),
+ fn=lambda img, rank: open_domain_classification(img, rank, return_all=False),
  inputs=[img_input, rank_dropdown],
  outputs=[open_domain_output],
  )
 
+
  zero_shot_btn.click(
  fn=zero_shot_classification,
  inputs=[img_input_zs, classes_txt],