Allows specifying top k results for custom class list (#31)

This change also changes the predictions to be returned sorted in
descending order for each image, like in the case of no custom classes.
Arguably this ordering is a more useful output whether asking for the
top k predictions or all of them.

Author: hlapp

src/bioclip/__main__.py

@@ -33,7 +33,7 @@ def predict(image_file: list[str], format: str,  output: str,
              cls_str: str, device: str,  rank: Rank, k: int):
     if cls_str:
         classifier = CustomLabelsClassifier(cls_ary=cls_str.split(','), device=device)
-        predictions = classifier.predict(image_paths=image_file)
+        predictions = classifier.predict(image_paths=image_file, k=k)
         write_results(predictions, format, output)
     else:
         classifier = TreeOfLifeClassifier(device=device)
@@ -87,8 +87,8 @@ def parse_args(input_args=None):
     if args.command == 'predict':
         if args.cls:
             # custom class list mode
-            if args.rank or args.k:
-                raise ValueError("Cannot use --cls with --rank or --k")
+            if args.rank:
+                raise ValueError("Cannot use --cls with --rank")
         else:
             # tree of life class list mode
             if not args.rank:

src/bioclip/predict.py

@@ -229,16 +229,20 @@ def _get_txt_features(self, classnames):
         return all_features
 
     @torch.no_grad()
-    def predict(self, image_paths: List[str] | str) -> dict[str, float]:
+    def predict(self, image_paths: List[str] | str, k: int = None) -> dict[str, float]:
         if isinstance(image_paths, str):
             image_paths = [image_paths]
         probs = self.create_probabilities_for_image_paths(image_paths, self.txt_features)
         result = []
         for image_path in image_paths:
-            for cls_str, prob in zip(self.classes, probs[image_path]):
+            img_probs = probs[image_path]
+            if not k or k > len(self.classes):
+                k = len(self.classes)
+            topk = img_probs.topk(k)
+            for i, prob in zip(topk.indices, topk.values):
                 result.append({
                     PRED_FILENAME_KEY: image_path,
-                    PRED_CLASSICATION_KEY: cls_str,
+                    PRED_CLASSICATION_KEY: self.classes[i],
                     PRED_SCORE_KEY: prob.item()
                 })
         return result

tests/test_main.py

@@ -44,9 +44,10 @@ def test_parse_args(self):
         # test error when using --cls with --rank
         with self.assertRaises(ValueError):
             parse_args(['predict', 'image.jpg', '--cls', 'class1,class2', '--rank', 'genus'])
-        # test error when using --cls with --k
-        with self.assertRaises(ValueError):
-            parse_args(['predict', 'image.jpg', '--cls', 'class1,class2', '--k', '10'])
+
+        # not an error when using --cls with --k
+        args = parse_args(['predict', 'image.jpg', '--cls', 'class1,class2', '--k', '10'])
+        self.assertEqual(args.k, 10)
 
         args = parse_args(['embed', 'image.jpg'])
         self.assertEqual(args.command, 'embed')