Merge pull request #8 from Imageomics/json-payload

Improve performance and prediction results

Author: johnbradley

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__':