Merge pull request #88 from Imageomics/87-subset-fix

Fix subsetting range bug

src/bioclip/predict.py

@@ -544,6 +544,8 @@ def apply_filter(self, keep_labels_ary: List[bool]):
         self._subset_txt_names = names
 
     def format_species_probs(self, image_key: str, probs: torch.Tensor, k: int = 5) -> List[dict[str, float]]:
+        # Prevent error when probs is smaller than k
+        k = min(k, probs.shape[0])
         topk = probs.topk(k)
         result = []
         for i, prob in zip(topk.indices, topk.values):

tests/test_predict.py

@@ -281,6 +281,19 @@ def test_create_taxa_filter_from_csv(self):
     def test_get_rank_labels(self):
         self.assertEqual(','.join(get_rank_labels()), 'kingdom,phylum,class,order,family,genus,species')
 
+    def test_format_species_probs_too_few_species(self):
+        classifier = TreeOfLifeClassifier()
+
+        # test when k < number of probabilities
+        probs = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7])
+        top_probs = classifier.format_species_probs(EXAMPLE_CAT_IMAGE, probs, k=5)
+        self.assertEqual(len(top_probs), 5)
+        self.assertEqual(top_probs[0]['file_name'], EXAMPLE_CAT_IMAGE)
+
+        # test when k > number of probabilities
+        probs = torch.tensor([0.1, 0.2, 0.3, 0.4])
+        top_probs = classifier.format_species_probs(EXAMPLE_CAT_IMAGE, probs, k=5)
+
 
 class TestEmbed(unittest.TestCase):
     def test_get_image_features(self):