Simplify binning classifier

src/bioclip/__main__.py

@@ -46,7 +46,8 @@ def predict(image_file: list[str],
         write_results(predictions, format, output)
     elif bins_path:
         bin_df = pd.read_csv(bins_path, index_col=0)
-        classifier = CustomLabelsBinningClassifier(bin_df=bin_df, **kwargs)
+        cls_to_bin = bin_df.bin.to_dict()
+        classifier = CustomLabelsBinningClassifier(cls_to_bin=cls_to_bin, **kwargs)
         predictions = classifier.predict(image_paths=image_file, k=k)
         write_results(predictions, format, output)
     else:

src/bioclip/predict.py

@@ -11,7 +11,6 @@
 from huggingface_hub import hf_hub_download
 from typing import Union, List
 from enum import Enum
-import pandas as pd
 
 
 HF_DATAFILE_REPO = "imageomics/bioclip-demo"
@@ -270,16 +269,15 @@ def group_probs(self, image_path: str, img_probs: torch.Tensor, k: int = None) -
 
 
 class CustomLabelsBinningClassifier(CustomLabelsClassifier):
-    def __init__(self, bin_df: pd.DataFrame, **kwargs):
-        classes = list(bin_df.index.values)
-        super().__init__(cls_ary=classes, **kwargs)
-        self.bin_df = bin_df
+    def __init__(self, cls_to_bin: dict, **kwargs):
+        super().__init__(cls_ary=cls_to_bin.keys(), **kwargs)
+        self.cls_to_bin = cls_to_bin
 
     def group_probs(self, image_path: str, img_probs: torch.Tensor, k: int = None) -> List[dict[str, float]]:
         result = []
         output = collections.defaultdict(float)
         for i in range(len(self.classes)):
-            name = self.bin_df.loc[self.classes[i]].bin
+            name = self.cls_to_bin[self.classes[i]]
             output[name] += img_probs[i]
         topk_names = heapq.nlargest(k, output, key=output.get)
         for name in topk_names: