Merge pull request #16 from Imageomics/20-txt-cache

Avoid recomputing text embeddings for custom labels

Author: johnbradley

README.md

@@ -89,8 +89,8 @@ df = pd.DataFrame(predictions)
 ```python
 from bioclip import CustomLabelsClassifier
 
-classifier = CustomLabelsClassifier()
-predictions = classifier.predict("Ursus-arctos.jpeg", ["duck","fish","bear"])
+classifier = CustomLabelsClassifier(["duck","fish","bear"])
+predictions = classifier.predict("Ursus-arctos.jpeg")
 for prediction in predictions:
    print(prediction["classification"], prediction["score"])
 ```

src/bioclip/__main__.py

@@ -31,10 +31,10 @@ def write_results_to_file(df, format, outfile):
 def predict(image_file: list[str], format: str,  output: str,
              cls_str: str, device: str,  rank: Rank, k: int):
     if cls_str:
-        classifier = CustomLabelsClassifier(device=device)
+        classifier = CustomLabelsClassifier(cls_ary=cls_str.split(','), device=device)
         data = []
         for image_path in image_file:
-            data.extend(classifier.predict(image_path=image_path, cls_ary=cls_str.split(',')))
+            data.extend(classifier.predict(image_path=image_path))
         write_results(data, format, output)
     else:
         classifier = TreeOfLifeClassifier(device=device)

src/bioclip/predict.py

@@ -166,13 +166,16 @@ def create_bioclip_tokenizer(tokenizer_str="ViT-B-16"):
 
 
 class CustomLabelsClassifier(object):
-    def __init__(self, device: Union[str, torch.device] = 'cpu', model_str: str = MODEL_STR):
+    def __init__(self, cls_ary: List[str], device: Union[str, torch.device] = 'cpu', model_str: str = MODEL_STR):
         self.device = device
         self.model = create_bioclip_model(device=device, model_str=model_str)
         self.model_str = model_str
         self.tokenizer = create_bioclip_tokenizer()
+        self.classes = [cls.strip() for cls in cls_ary]
+        self.txt_features = self._get_txt_features(self.classes)
 
-    def get_txt_features(self, classnames):
+    @torch.no_grad()
+    def _get_txt_features(self, classnames):
         all_features = []
         for classname in classnames:
             txts = [template(classname) for template in OPENA_AI_IMAGENET_TEMPLATE]
@@ -185,19 +188,17 @@ def get_txt_features(self, classnames):
         return all_features
 
     @torch.no_grad()
-    def predict(self, image_path: str, cls_ary: List[str]) -> dict[str, float]:
+    def predict(self, image_path: str) -> dict[str, float]:
         img = open_image(image_path)
-        classes = [cls.strip() for cls in cls_ary]
-        txt_features = self.get_txt_features(classes)
 
         img = preprocess_img(img).to(self.device)
         img_features = self.model.encode_image(img.unsqueeze(0))
         img_features = F.normalize(img_features, dim=-1)
 
-        logits = (self.model.logit_scale.exp() * img_features @ txt_features).squeeze()
+        logits = (self.model.logit_scale.exp() * img_features @ self.txt_features).squeeze()
         probs = F.softmax(logits, dim=0).to("cpu").tolist()
         pred_list = []
-        for cls, prob in zip(classes, probs):
+        for cls, prob in zip(self.classes, probs):
             pred_list.append({
                 PRED_FILENAME_KEY: image_path,
                 PRED_CLASSICATION_KEY: cls,
@@ -207,8 +208,8 @@ def predict(self, image_path: str, cls_ary: List[str]) -> dict[str, float]:
 
 
 def predict_classifications_from_list(img: Union[PIL.Image.Image, str], cls_ary: List[str], device: Union[str, torch.device] = 'cpu') -> dict[str, float]:
-    classifier = CustomLabelsClassifier(device=device)
-    return classifier.predict(img, cls_ary)
+    classifier = CustomLabelsClassifier(cls_ary=cls_ary, device=device)
+    return classifier.predict(img)
 
 
 def get_tol_classification_labels(rank: Rank) -> List[str]:
@@ -293,7 +294,6 @@ def format_grouped_probs(self, image_path: str, probs: torch.Tensor, rank: Rank,
         for name in topk_names:
             item = { PRED_FILENAME_KEY: image_path }
             item.update(name_to_class_dict[name])
-            #item.update(class_dict_lookup)
             item[PRED_SCORE_KEY] = output[name].item()
             prediction_ary.append(item)
         return prediction_ary

tests/test_predict.py

@@ -45,8 +45,8 @@ def test_tree_of_life_classifier_family(self):
         self.assertEqual(prediction_ary[0], prediction_dict)
 
     def test_custom_labels_classifier(self):
-        classifier = CustomLabelsClassifier()
-        results = classifier.predict(image_path=EXAMPLE_CAT_IMAGE, cls_ary=['cat', 'dog'])
+        classifier = CustomLabelsClassifier(cls_ary=['cat', 'dog'])
+        results = classifier.predict(image_path=EXAMPLE_CAT_IMAGE)
         self.assertEqual(results, [
             {'file_name': EXAMPLE_CAT_IMAGE, 'classification': 'cat', 'score': unittest.mock.ANY},
             {'file_name': EXAMPLE_CAT_IMAGE, 'classification': 'dog', 'score': unittest.mock.ANY},