Merge pull request #37 from Imageomics/30-classes-file

Add --cls-file to predict command

Author: johnbradley

README.md

@@ -119,7 +119,7 @@ Options:
   --format=FORMAT      format of the output (table or csv) for predict mode [default: csv]
   --rank=RANK          rank of the classification (kingdom, phylum, class, order, family, genus, species) [default: species] 
   --k=K                number of top predictions to show [default: 5]
-  --cls=CLS            comma separated list of classes to predict, when specified the --rank and --k arguments are not allowed
+  --cls=CLS            classes to predict: either a comma separated list or a path to a text file of classes (one per line), when specified the --rank argument is not allowed.
   --device=DEVICE      device to use matrix math (cpu or cuda or mps) [default: cpu]
   --output=OUTFILE     print output to file OUTFILE [default: stdout]
 ```

src/bioclip/__main__.py

@@ -1,6 +1,7 @@
 from bioclip import TreeOfLifeClassifier, Rank, CustomLabelsClassifier
 from .predict import BIOCLIP_MODEL_STR
 import open_clip as oc
+import os
 import json
 import sys
 import prettytable as pt
@@ -83,7 +84,9 @@ def create_parser():
     predict_parser.add_argument('--rank', choices=['kingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'],
                                 help='rank of the classification, default: species (when)')
     predict_parser.add_argument('--k', type=int, help='number of top predictions to show, default: 5')
-    predict_parser.add_argument('--cls', help='comma separated list of classes to predict, when specified the --rank argument is not allowed')
+    cls_help = "classes to predict: either a comma separated list or a path to a text file of classes (one per line), when specified the --rank argument is not allowed."
+    predict_parser.add_argument('--cls', help=cls_help)
+
     predict_parser.add_argument('--device', **device_arg)
     predict_parser.add_argument('--model', **model_arg)
     predict_parser.add_argument('--pretrained', **pretrained_arg)
@@ -128,6 +131,13 @@ def parse_args(input_args=None):
     return args
 
 
+def create_classes_str(cls_file_path):
+    """Reads a file with one class per line and returns a comma separated string of classes"""
+    with open(cls_file_path, 'r') as cls_file:
+        cls_str = [item.strip() for item in cls_file.readlines()]
+    return ",".join(cls_str)
+
+
 def main():
     args = parse_args()
     if args.command == 'embed':
@@ -137,10 +147,13 @@ def main():
               model_str=args.model,
               pretrained_str=args.pretrained)
     elif args.command == 'predict':
+        cls_str = args.cls
+        if args.cls and os.path.exists(args.cls):
+            cls_str = create_classes_str(args.cls)
         predict(args.image_file,
                 format=args.format,
                 output=args.output,
-                cls_str=args.cls,
+                cls_str=cls_str,
                 rank=args.rank,
                 k=args.k,
                 device=args.device,

tests/test_main.py

@@ -1,5 +1,7 @@
 import unittest
-from bioclip.__main__ import parse_args, Rank
+from unittest.mock import mock_open, patch
+import argparse
+from bioclip.__main__ import parse_args, Rank, create_classes_str, main
 
 
 class TestParser(unittest.TestCase):
@@ -49,6 +51,10 @@ def test_parse_args(self):
         args = parse_args(['predict', 'image.jpg', '--cls', 'class1,class2', '--k', '10'])
         self.assertEqual(args.k, 10)
 
+        # example showing filename
+        args = parse_args(['predict', 'image.jpg', '--cls', 'classes.txt', '--k', '10'])
+        self.assertEqual(args.cls, 'classes.txt')
+
         args = parse_args(['embed', 'image.jpg'])
         self.assertEqual(args.command, 'embed')
         self.assertEqual(args.image_file, ['image.jpg'])
@@ -60,3 +66,43 @@ def test_parse_args(self):
         self.assertEqual(args.image_file, ['image.jpg', 'image2.png'])
         self.assertEqual(args.output, 'data.json')
         self.assertEqual(args.device, 'cuda')
+
+    def test_create_classes_str(self):
+        data = "class1\nclass2\nclass3"
+        with patch("builtins.open", mock_open(read_data=data)) as mock_file:
+            self.assertEqual(create_classes_str('path/to/file'), 'class1,class2,class3')
+
+    @patch('bioclip.__main__.predict')
+    @patch('bioclip.__main__.parse_args')
+    def test_predict_no_class(self, mock_parse_args, mock_predict):
+        mock_parse_args.return_value = argparse.Namespace(command='predict', image_file='image.jpg', format='csv',
+                                                          output='stdout', rank=Rank.SPECIES, k=5, cls=None, device='cpu',
+                                                          model=None, pretrained=None)
+        main()
+        mock_predict.assert_called_with('image.jpg', format='csv', output='stdout', cls_str=None, rank=Rank.SPECIES, k=5,
+                                        device='cpu', model_str=None, pretrained_str=None)
+
+    @patch('bioclip.__main__.predict')
+    @patch('bioclip.__main__.parse_args')
+    @patch('bioclip.__main__.os')
+    def test_predict_class_list(self, mock_os, mock_parse_args, mock_predict):
+        mock_os.path.exists.return_value = False
+        mock_parse_args.return_value = argparse.Namespace(command='predict', image_file='image.jpg', format='csv',
+                                                          output='stdout', rank=Rank.SPECIES, k=5, cls='dog,fish,bird',
+                                                          device='cpu', model=None, pretrained=None)
+        main()
+        mock_predict.assert_called_with('image.jpg', format='csv', output='stdout', cls_str='dog,fish,bird', rank=Rank.SPECIES,
+                                        k=5, device='cpu', model_str=None, pretrained_str=None)
+
+    @patch('bioclip.__main__.predict')
+    @patch('bioclip.__main__.parse_args')
+    @patch('bioclip.__main__.os')
+    def test_predict_class_file(self, mock_os, mock_parse_args, mock_predict):
+        mock_os.path.exists.return_value = True
+        mock_parse_args.return_value = argparse.Namespace(command='predict', image_file='image.jpg', format='csv', 
+                                                          output='stdout', rank=Rank.SPECIES, k=5, cls='somefile.txt',
+                                                          device='cpu', model=None, pretrained=None)
+        with patch("builtins.open", mock_open(read_data='dog\nfish\nbird')) as mock_file:
+            main()
+        mock_predict.assert_called_with('image.jpg', format='csv', output='stdout', cls_str='dog,fish,bird', rank=Rank.SPECIES,
+                                        k=5, device='cpu', model_str=None, pretrained_str=None)