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