add new method api_classification

app.py

@@ -153,6 +153,41 @@ def open_domain_classification(img, rank: int) -> dict[str, float]:
     return {name: output[name] for name in topk_names}
 
 
+
+@torch.no_grad()
+def api_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]
+
+    logger.info(">>>>")
+    logger.info(probs[0])
+
+    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)