Improve performance and prediction results

Introduces two classes that will cache the model so creating
predictions for multiple images will only load the model and the
tree of life embeddings once.

Changes the prediction returned values to be more consistent.
One deviation from what was requested with issue #1 is flattening
the results to simplify creating a pandas DataFrame.

Fixes #1

Simplifies the command line interface to default to csv format and process
multiple images.

Adds tests that create predictions. These test require downloading the model
so may be inappropriate for automated testing.
Tests can be run with pytest.

Co-authored-by: Hilmar Lapp <[email protected]>
Co-authored-by: Elizabeth Campolongo <[email protected]>

Author: 3 people

README.md

@@ -12,9 +12,8 @@ Command line tool and python package to simplify using [BioCLIP](https://imageom
 **Table of Contents**
 
 - [Installation](#installation)
-- [Command Line Usage](#command-line-usage)
 - [Python Package Usage](#python-package-usage)
-- [License](#license)
+- [Command Line Usage](#command-line-usage)
 - [Acknowledgments](#acknowledgments)
 - [License](#license)
 
@@ -29,132 +28,162 @@ pip install git+https://github.com/Imageomics/pybioclip
 
 If you have any issues with installation, please first upgrade pip by running `pip install --upgrade pip`.
 
-## Command Line Usage
+## Python Package Usage
+### Predict species classification
 
-### Predict classification
+```python
+from bioclip import TreeOfLifeClassifier, Rank
 
-#### Example: Predict species for an image
-The example image used below is [`Ursus-arctos.jpeg`](https://huggingface.co/spaces/imageomics/bioclip-demo/blob/ef075807a55687b320427196ac1662b9383f988f/examples/Ursus-arctos.jpeg) from the [bioclip-demo](https://huggingface.co/spaces/imageomics/bioclip-demo).
+classifier = TreeOfLifeClassifier()
+predictions = classifier.predict("Ursus-arctos.jpeg", Rank.SPECIES)
 
-Predict species for an `Ursus-arctos.jpeg` file:
-```console
-bioclip predict Ursus-arctos.jpeg
-```
-Output:
+for prediction in predictions:
+    print(prediction["species"], "-", prediction["score"])
 ```
-+----------------------------------------------------------------------------------------+-----------------------+
-|                                         Taxon                                          |      Probability      |
-+----------------------------------------------------------------------------------------+-----------------------+
-|        Animalia Chordata Mammalia Carnivora Ursidae Ursus arctos (Kodiak bear)         |   0.9356034994125366  |
-| Animalia Chordata Mammalia Carnivora Ursidae Ursus arctos syriacus (syrian brown bear) |  0.05616999790072441  |
-|          Animalia Chordata Mammalia Carnivora Ursidae Ursus arctos bruinosus           |  0.004126196261495352 |
-|               Animalia Chordata Mammalia Carnivora Ursidae Ursus arctus                | 0.0024959812872111797 |
-|  Animalia Chordata Mammalia Carnivora Ursidae Ursus americanus (Louisiana black bear)  | 0.0005009894957765937 |
-+----------------------------------------------------------------------------------------+-----------------------+
-```
-
----
 
-To save as a CSV or JSON file you can use the `--format <file type>` and `--output <filename>` arguments with `csv` or `json`, respectively.
-
-To save the JSON output to `ursus.json` run:
+Output:
 ```console
-bioclip predict --format json --output ursus.json Ursus-arctos.jpeg 
-```
+Ursus arctos - 0.9356034994125366
+Ursus arctos syriacus - 0.05616999790072441
+Ursus arctos bruinosus - 0.004126196261495352
+Ursus arctus - 0.0024959812872111797
+Ursus americanus - 0.0005009894957765937
+```
+
+Output from the `predict()` method showing the dictionary structure:
+```
+[{
+    'kingdom': 'Animalia',
+    'phylum': 'Chordata',
+    'class': 'Mammalia',
+    'order': 'Carnivora',
+    'family': 'Ursidae',
+    'genus': 'Ursus',
+    'species_epithet': 'arctos',
+    'species': 'Ursus arctos',
+    'common_name': 'Kodiak bear'
+    'score': 0.9356034994125366
+}]
+```
+
+The output from the predict function can be converted into a [pandas DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html) like so:
+```python
+import pandas as pd
+from bioclip import TreeOfLifeClassifier, Rank
 
-To save the CSV output to `ursus.csv` run:
-```console
-bioclip predict --format csv --output ursus.csv Ursus-arctos.jpeg 
+classifier = TreeOfLifeClassifier()
+predictions = classifier.predict("Ursus-arctos.jpeg", Rank.SPECIES)
+df = pd.DataFrame(predictions)
 ```
 
-#### Predict genus for an image
+### Predict from a list of classes
+```python
+from bioclip import CustomLabelsClassifier
 
-Predict genus for image `Ursus-arctos.jpeg`, restricted to the top 3 predictions:
-```console
-bioclip predict --rank genus --k 3 Ursus-arctos.jpeg
+classifier = CustomLabelsClassifier()
+predictions = classifier.predict("Ursus-arctos.jpeg", ["duck","fish","bear"])
+for prediction in predictions:
+   print(prediction["classification"], prediction["score"])
 ```
 Output:
+```console
+duck 1.0306726583309e-09
+fish 2.932403668845507e-12
+bear 1.0
 ```
-+---------------------------------------------------------+------------------------+
-|                          Taxon                          |      Probability       |
-+---------------------------------------------------------+------------------------+
-|    Animalia Chordata Mammalia Carnivora Ursidae Ursus   |   0.9994320273399353   |
-| Animalia Chordata Mammalia Artiodactyla Cervidae Cervus | 0.00032594642834737897 |
-|  Animalia Chordata Mammalia Artiodactyla Cervidae Alces | 7.803700282238424e-05  |
-+---------------------------------------------------------+------------------------+
+
+## Command Line Usage
 ```
+bioclip predict [options] [IMAGE_FILE...]
 
-#### Optional arguments for predicting classifications:
-- `--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]
-- `--format FORMAT` - format of the output (table, json, or csv) [default: table]
-- `--output OUTPUT` - save output to a filename instead of printing it [default: stdout]
+Arguments:
+  IMAGE_FILE         input image file
 
+Options:
+  -h --help
+  --format=FORMAT    format of the output (table or csv) [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 ignored [default: all]
+  --output=OUTFILE   print output to file OUTFILE [default: stdout]
+```
 
-### Predict from a list of classes
+### Predict classification
 
-Create predictions for 3 classes (cat, bird, and bear) for image `Ursus-arctos.jpeg`:
+#### Predict species for an image
+The example images used below are [`Ursus-arctos.jpeg`](https://huggingface.co/spaces/imageomics/bioclip-demo/blob/ef075807a55687b320427196ac1662b9383f988f/examples/Ursus-arctos.jpeg) 
+and [`Felis-catus.jpeg`](https://huggingface.co/spaces/imageomics/bioclip-demo/blob/ef075807a55687b320427196ac1662b9383f988f/examples/Felis-catus.jpeg) both from the [bioclip-demo](https://huggingface.co/spaces/imageomics/bioclip-demo).
+
+Predict species for an `Ursus-arctos.jpeg` file:
 ```console
-bioclip predict --cls cat,bird,bear Ursus-arctos.jpeg
+bioclip predict Ursus-arctos.jpeg
 ```
 Output:
 ```
-+-------+-----------------------+
-| Taxon |      Probability      |
-+-------+-----------------------+
-|  cat  | 4.581644930112816e-08 |
-|  bird | 3.051998476166773e-08 |
-|  bear |   0.9999998807907104  |
-+-------+-----------------------+%                                                                  
+bioclip predict Ursus-arctos.jpeg
+file_name,kingdom,phylum,class,order,family,genus,species_epithet,species,common_name,score
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,arctos,Ursus arctos,Kodiak bear,0.9356034994125366
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,arctos syriacus,Ursus arctos syriacus,syrian brown bear,0.05616999790072441
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,arctos bruinosus,Ursus arctos bruinosus,,0.004126196261495352
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,arctus,Ursus arctus,,0.0024959812872111797
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,americanus,Ursus americanus,Louisiana black bear,0.0005009894957765937
 ```
 
-#### Optional arguments for predicting from a list of classes:
-- `--format FORMAT` - format of the output (table, json, or csv) [default: table]
-- `--output OUTPUT` - save output to a filename instead of printing it [default: stdout]
-- `--cls CLS` - comma separated list of classes to predict, when specified the `--rank` and `--k` arguments are ignored [default: all]
+#### Predict species for multiple images saving to a file
 
-
-### View command line help
+To make predictions for files `Ursus-arctos.jpeg` and `Felis-catus.jpeg` saving the output to a file named `predictions.csv`:
 ```console
-bioclip --help
+bioclip predict --output predictions.csv Ursus-arctos.jpeg Felis-catus.jpeg
+```
+The contents of `predictions.csv` will look like this: 
+```
+file_name,kingdom,phylum,class,order,family,genus,species_epithet,species,common_name,score
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,arctos,Ursus arctos,Kodiak bear,0.9356034994125366
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,arctos syriacus,Ursus arctos syriacus,syrian brown bear,0.05616999790072441
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,arctos bruinosus,Ursus arctos bruinosus,,0.004126196261495352
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,arctus,Ursus arctus,,0.0024959812872111797
+Ursus-arctos.jpeg,Animalia,Chordata,Mammalia,Carnivora,Ursidae,Ursus,americanus,Ursus americanus,Louisiana black bear,0.0005009894957765937
+Felis-catus.jpeg,Animalia,Chordata,Mammalia,Carnivora,Felidae,Felis,silvestris,Felis silvestris,European Wildcat,0.7221033573150635
+Felis-catus.jpeg,Animalia,Chordata,Mammalia,Carnivora,Felidae,Felis,catus,Felis catus,Domestic Cat,0.19810837507247925
+Felis-catus.jpeg,Animalia,Chordata,Mammalia,Carnivora,Felidae,Felis,margarita,Felis margarita,Sand Cat,0.02798456884920597
+Felis-catus.jpeg,Animalia,Chordata,Mammalia,Carnivora,Felidae,Lynx,felis,Lynx felis,,0.021829601377248764
+Felis-catus.jpeg,Animalia,Chordata,Mammalia,Carnivora,Felidae,Felis,bieti,Felis bieti,Chinese desert cat,0.010979168117046356
 ```
 
-## Python Package Usage
-### Predict species classification
-
-```python
-from bioclip import predict_classification, Rank
-
-predictions = predict_classification("Ursus-arctos.jpeg", Rank.SPECIES)
-
-for species_name, probability in predictions.items():
-   print(species_name, probability)
+#### Predict top 3 genera for an image and display output as a table
+```console
+bioclip predict --format table --k 3 --rank=genus Ursus-arctos.jpeg
 ```
 
 Output:
-```console
-Animalia Chordata Mammalia Carnivora Ursidae Ursus arctos (Kodiak bear) 0.9356034994125366
-Animalia Chordata Mammalia Carnivora Ursidae Ursus arctos syriacus (syrian brown bear) 0.05616999790072441
-Animalia Chordata Mammalia Carnivora Ursidae Ursus arctos bruinosus 0.004126196261495352
-Animalia Chordata Mammalia Carnivora Ursidae Ursus arctus 0.0024959812872111797
-Animalia Chordata Mammalia Carnivora Ursidae Ursus americanus (Louisiana black bear) 0.0005009894957765937
+```
++-------------------+----------+----------+----------+--------------+----------+--------+------------------------+
+|     file_name     | kingdom  |  phylum  |  class   |    order     |  family  | genus  |         score          |
++-------------------+----------+----------+----------+--------------+----------+--------+------------------------+
+| Ursus-arctos.jpeg | Animalia | Chordata | Mammalia |  Carnivora   | Ursidae  | Ursus  |   0.9994320273399353   |
+| Ursus-arctos.jpeg | Animalia | Chordata | Mammalia | Artiodactyla | Cervidae | Cervus | 0.00032594642834737897 |
+| Ursus-arctos.jpeg | Animalia | Chordata | Mammalia | Artiodactyla | Cervidae | Alces  | 7.803700282238424e-05  |
++-------------------+----------+----------+----------+--------------+----------+--------+------------------------+
 ```
 
 ### Predict from a list of classes
-```python
-from bioclip import predict_classifications_from_list, Rank
-
-predictions = predict_classifications_from_list("Ursus-arctos.jpeg",
-                                                ["duck","fish","bear"])
-
-for cls, probability in predictions.items():
-   print(cls, probability)
+Create predictions for 3 classes (cat, bird, and bear) for image `Ursus-arctos.jpeg`:
+```console
+bioclip predict --cls cat,bird,bear Ursus-arctos.jpeg
 ```
 Output:
+```
+file_name,classification,score
+Ursus-arctos.jpeg,cat,4.581644930112816e-08
+Ursus-arctos.jpeg,bird,3.051998476166773e-08
+Ursus-arctos.jpeg,bear,0.9999998807907104                                                                 
+```
+
+### View command line help
 ```console
-duck 1.0306726583309e-09
-fish 2.932403668845507e-12
-bear 1.0
+bioclip --help
 ```
 
 ## License

pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
   'torch',
   'docopt-ng',
   'prettytable',
+  'pandas',
 ]
 
 [project.urls]
@@ -90,3 +91,8 @@ exclude_lines = [
   "if __name__ == .__main__.:",
   "if TYPE_CHECKING:",
 ]
+
+[tool.pytest.ini_options]
+pythonpath = [
+  "src"
+]

src/bioclip/__init__.py

@@ -1,6 +1,6 @@
 # SPDX-FileCopyrightText: 2024-present John Bradley <[email protected]>
 #
 # SPDX-License-Identifier: MIT
-from bioclip.predict import predict_classification, Rank, predict_classifications_from_list
+from bioclip.predict import TreeOfLifeClassifier, Rank, CustomLabelsClassifier
 
-__all__ = ["predict_classification", "Rank", "predict_classifications_from_list"]
+__all__ = ["TreeOfLifeClassifier", "Rank", "CustomLabelsClassifier"]

src/bioclip/__main__.py

@@ -1,4 +1,4 @@
-"""Usage: bioclip predict [options] IMAGE_FILE 
+"""Usage: bioclip predict [options] [IMAGE_FILE...]
 
 Use BioCLIP to generate predictions for an IMAGE_FILE.
 
@@ -7,36 +7,41 @@
 
 Options:
   -h --help
-  --format=FORMAT      format of the output (table, json, or csv) [default: table]
+  --format=FORMAT      format of the output (table or csv) [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 ignored [default: all]
-  --output=OUTFILE     save output to a filename instead of printing it [default: stdout]
+  --output=OUTFILE     print output to file OUTFILE [default: stdout]
 
 """
 from docopt import docopt
-from bioclip import predict_classification, predict_classifications_from_list, Rank
+from bioclip import TreeOfLifeClassifier, Rank, CustomLabelsClassifier
 import json
 import sys
 import prettytable as pt
 import csv
+import pandas as pd
 
 
-def write_results(result, format, outfile):
+def write_results(data, format, output):
+    df = pd.DataFrame(data)
+    if output == 'stdout':
+        write_results_to_file(df, format, sys.stdout)
+    else:
+        with open(output, 'w') as outfile:
+            write_results_to_file(df, format, outfile)
+
+
+def write_results_to_file(df, format, outfile):
     if format == 'table':
         table = pt.PrettyTable()
-        table.field_names = ['Taxon', 'Probability']
-        for taxon, prob in result.items():
-            table.add_row([taxon, prob])
+        table.field_names = df.columns
+        for index, row in df.iterrows():
+            table.add_row(row)
         outfile.write(str(table))
         outfile.write('\n')
-    elif format == 'json':
-        json.dump(result, outfile, indent=2)
     elif format == 'csv':
-        writer = csv.writer(outfile)
-        writer.writerow(['Taxon', 'Probability'])
-        for taxon, prob in result.items():
-            writer.writerow([taxon, prob])
+        df.to_csv(outfile, index=False)
     else:
         raise ValueError(f"Invalid format: {format}")
 
@@ -48,22 +53,21 @@ def main():
     output = x['--output']
     image_file = x['IMAGE_FILE']
     cls = x['--cls']
-    if not format in ['table', 'json', 'csv']:
+    if not format in ['table', 'csv']:
         raise ValueError(f"Invalid format: {format}")
     rank = Rank[x['--rank'].upper()]
     if cls == 'all':
-        result = predict_classification(img=image_file,
-                                        rank=rank,
-                                        k=int(x['--k']))
-    else:
-        result = predict_classifications_from_list(img=image_file, 
-                                                   cls_ary=cls.split(','))
-    outfile = sys.stdout
-    if output == 'stdout':
-        write_results(result, format, sys.stdout)
+        classifier = TreeOfLifeClassifier()
+        data = []
+        for image_path in image_file:
+            data.extend(classifier.predict(image_path=image_path, rank=rank, k=int(x['--k'])))
+        write_results(data, format, output)
     else:
-        with open(output, 'w') as outfile:
-            write_results(result, format, outfile)
+        classifier = CustomLabelsClassifier()
+        data = []
+        for image_path in image_file:
+            data.extend(classifier.predict(image_path=image_path, cls_ary=cls.split(',')))
+        write_results(data, format, output)
 
 
 if __name__ == '__main__':