Repository: davanstrien/flyswot

Issue:
add support for multi-head models
- [ ] Add new dataclass or modify to support image prediction with multiple predictions
- [ ] Add support for writing to CSV file
- [ ] Modify table print out to print two tables 



Discussion:

Resolving patch:
diff --git a/src/flyswot/inference.py b/src/flyswot/inference.py
--- a/src/flyswot/inference.py
+++ b/src/flyswot/inference.py
@@ -4,20 +4,26 @@
 import time
 from abc import ABC
 from abc import abstractmethod
+from collections import defaultdict
+from collections import OrderedDict
 from dataclasses import asdict
 from dataclasses import dataclass
 from datetime import datetime
 from datetime import timedelta
+from functools import singledispatch
 from pathlib import Path
 from typing import Iterable
-from typing import Iterator
 from typing import List
+from typing import Tuple
 from typing import Union
 
 import numpy as np
 import onnxruntime as rt  # type: ignore
 import typer
 from PIL import Image  # type: ignore
+from rich import print
+from rich.columns import Columns
+from rich.layout import Layout
 from rich.table import Table
 from toolz import itertoolz
 
@@ -35,7 +41,7 @@ class ImagePredictionItem:
     Attributes:
         path: The Path to the image
         predicted_label: The predicted label i.e. the argmax value for the prediction tensor
-        condidence: The confidence for `predicted_label` i.e. the max value for prediction tensor
+        confidence: The confidence for `predicted_label` i.e. the max value for prediction tensor
     """
 
     path: Path
@@ -50,6 +56,14 @@ def __post_init__(self) -> Union[Path, None]:
             pass
 
 
+@dataclass
+class MultiLabelImagePredictionItem:
+    """Multiple predictions for a single image"""
+
+    path: Path
+    predictions: List[Tuple[str, float]]
+
+
 @dataclass
 class PredictionBatch:
     """Container for ImagePredictionItems"""
@@ -57,8 +71,22 @@ class PredictionBatch:
     batch: List[ImagePredictionItem]
 
     def __post_init__(self):
-        """Returns a list of all predicted labels in batch"""
-        self.batch_labels: Iterator[str] = (item.predicted_label for item in self.batch)
+        """Returns a iterable of all predicted labels in batch"""
+        self.batch_labels: Iterable[str] = (item.predicted_label for item in self.batch)
+
+
+@dataclass
+class MultiPredictionBatch:
+    """Container for MultiLabelImagePRedictionItems"""
+
+    batch: List[MultiLabelImagePredictionItem]
+
+    def __post_init__(self):
+        """Returns a iterable of lists containing all predicted labels in batch"""
+        self.batch_labels: Iterable = (
+            list(itertoolz.pluck(0, pred))
+            for pred in zip(*[o.predictions for o in self.batch])
+        )
 
 
 image_extensions = {k for k, v in mimetypes.types_map.items() if v.startswith("image/")}
@@ -72,6 +100,15 @@ def predict_image(
     pass  # pragma: no cover
 
 
+def check_files(files: List, pattern: str, directory: Path) -> None:
+    """Check if files exist and raises error if not"""
+    if not files:
+        typer.echo(
+            f"Didn't find any files maching {pattern} in {directory}, please check the inputs to flyswot"
+        )
+        raise typer.Exit(code=1)
+
+
 @app.command(name="directory")
 def predict_directory(
     directory: Path = typer.Argument(
@@ -91,7 +128,9 @@ def predict_directory(
     image_format: str = typer.Option(
         ".tif", help="Image format for flyswot to use for predictions"
     ),
-    check_latest: bool = typer.Option(True, help="Use latest available model"),
+    # check_latest: bool = typer.Option(True, help="Use latest available model"),
+    model_name: str = typer.Option("latest", help="Which model to use"),
+    model_path: str = None,
 ):
     """Predicts against all images stored under DIRECTORY which match PATTERN in the filename.
 
@@ -101,39 +140,57 @@ def predict_directory(
     """
     start_time = time.perf_counter()
     model_dir = models.ensure_model_dir()
-    # TODO add load learner function that can be passed a model name
-    model_parts = models.ensure_model(model_dir, check_latest)
-    model = model_parts.model
-    vocab = models.load_vocab(model_parts.vocab)
-    onnxinference = OnnxInferenceSession(model, vocab)
+    if model_name == "latest":
+        model_parts = models.ensure_model(model_dir, check_latest=True)
+    if model_name != "latest" and not model_path:
+        model_parts = models._get_model_parts(Path(model_dir / Path(model_name)))
+    if model_name != "latest" and model_path:
+        model_parts = models._get_model_parts(Path(model_path) / Path(model_name))
+    onnxinference = OnnxInferenceSession(model_parts.model, model_parts.vocab)
     files = list(core.get_image_files_from_pattern(directory, pattern, image_format))
-    if not files:
-        typer.echo(
-            f"Didn't find any files maching {pattern} in {directory}, please check the inputs to flyswot"
-        )
-        raise typer.Exit(code=1)
+    check_files(files, pattern, directory)
     typer.echo(f"Found {len(files)} files matching {pattern} in {directory}")
     csv_fname = create_csv_fname(csv_save_dir)
-    create_csv_header(csv_fname)
     with typer.progressbar(length=len(files)) as progress:
-        all_preds = []
-        predictions = []
-        for batch in itertoolz.partition_all(bs, files):
+        for i, batch in enumerate(itertoolz.partition_all(bs, files)):
             batch_predictions = onnxinference.predict_batch(batch, bs)
-            all_preds.append(batch_predictions.batch_labels)
-            predictions.append(batch_predictions)
+            if i == 0:
+                create_csv_header(batch_predictions, csv_fname)
+            write_batch_preds_to_csv(batch_predictions, csv_fname)
             progress.update(len(batch))
-            write_batch_preds_to_csv(csv_fname, batch_predictions)
-        all_preds = list(itertoolz.concat(all_preds))
     typer.echo(f"CSV report stored in {csv_fname}")
     delta = timedelta(seconds=time.perf_counter() - start_time)
     typer.echo(f"Time taken to run:  {str(delta)}")
-    print_table(all_preds)
-
-
-def print_table(decoded) -> None:
+    print_inference_summary(csv_fname)
+
+
+def print_inference_summary(csv_fname):
+    """print_inference_summary from `fname`"""
+    labels_to_print = labels_from_csv(csv_fname)
+    tables = [
+        print_table(labels, f"Prediction summary {i+1}", print=False)
+        for i, labels in enumerate(labels_to_print)
+    ]
+    columns = Columns(tables)
+    print(columns)
+
+
+def labels_from_csv(fname: Path) -> List[List[str]]:
+    """Gets labels from csv `fname`"""
+    columns = defaultdict(list)
+    with open(fname, "r") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            for (k, v) in row.items():
+                columns[k].append(v)
+    return [columns[k] for k in columns if "prediction" in k]
+
+
+def print_table(
+    decoded: list, header: str = "Prediction summary", print: bool = True
+) -> Table:
     """Prints table summary of predicted labels"""
-    table = Table(show_header=True, title="Prediction summary")
+    table = Table(show_header=True, title=header)
     table.add_column(
         "Class",
     )
@@ -150,7 +207,33 @@ def print_table(decoded) -> None:
             table.add_row("Total", str(total), "")
         else:
             table.add_row(key, str(count), f"{percentage}")
-    console.print(table)
+    if print:
+        console.print(table)
+    return table
+
+
+def make_layout():
+    """Define the layout."""
+    layout = Layout(name="root")
+    layout.split(Layout(name="header", size=4), Layout(name="main"))
+    layout["main"].split_column(
+        Layout(name="info", size=4), Layout(name="body", ratio=2, minimum_size=60)
+    )
+    layout["info"].split_row(Layout(name="time"), Layout(name="files"))
+    return layout
+
+
+# def print_summary(columns, time, files):
+#     layout = make_layout()
+#     MARKDOWN = """
+#     # Prediction summary
+#     """
+#     md = Markdown(MARKDOWN)
+#     layout["header"].update(md)
+#     layout["body"].update(columns)
+#     layout["time"].update(time)
+#     layout["files"].update(f"Number of files checked {len(files)}")
+#     print(layout)
 
 
 def create_csv_fname(csv_directory: Path) -> Path:
@@ -161,15 +244,44 @@ def create_csv_fname(csv_directory: Path) -> Path:
     return Path(csv_directory / fname)
 
 
-def create_csv_header(csv_path: Path) -> None:
+@singledispatch
+def create_csv_header(batch, csv_path) -> None:
+    """Print csv header"""
+    raise NotImplementedError(f"Not implemented for type {batch}")
+
+
+@create_csv_header.register
+def _(batch: PredictionBatch, csv_path):
     """Creates a header for csv `csv_path`"""
     with open(csv_path, mode="w", newline="") as csv_file:
-        field_names = ["path", "directory", "predicted_label", "confidence"]
+        field_names = ["path", "directory", "prediction", "confidence"]
+        writer = csv.DictWriter(csv_file, fieldnames=field_names)
+        writer.writeheader()
+
+
+@create_csv_header.register
+def _(batch: MultiPredictionBatch, csv_path):
+    item = batch.batch[0]
+    pred = OrderedDict()
+    pred["path"] = asdict(item)["path"]
+    pred["directory"] = asdict(item)["path"].parent
+    for k, v in enumerate(item.predictions):
+        pred[f"prediction_label_{k}"] = v[0]
+        pred[f"confidence__label_{k}"] = v[1]
+    with open(csv_path, mode="w", newline="") as csv_file:
+        field_names = list(pred.keys())
         writer = csv.DictWriter(csv_file, fieldnames=field_names)
         writer.writeheader()
 
 
-def write_batch_preds_to_csv(csv_fpath: Path, predictions: PredictionBatch) -> None:
+@singledispatch
+def write_batch_preds_to_csv(predictions, csv_fpath: Path) -> None:
+    """print batch preds to csv"""
+    raise NotImplementedError(f"Not implemented for type {predictions}")
+
+
+@write_batch_preds_to_csv.register
+def _(predictions: PredictionBatch, csv_fpath: Path) -> None:
     """Appends `predictions` batch to `csv_path`"""
     with open(csv_fpath, mode="a", newline="") as csv_file:
         field_names = ["path", "directory", "predicted_label", "confidence"]
@@ -180,22 +292,39 @@ def write_batch_preds_to_csv(csv_fpath: Path, predictions: PredictionBatch) -> N
             writer.writerow(row)
 
 
+@write_batch_preds_to_csv.register
+def _(predictions: MultiPredictionBatch, csv_fpath) -> None:
+    for pred in predictions.batch:
+        row = OrderedDict()
+        row["path"] = asdict(pred)["path"]
+        row["directory"] = asdict(pred)["path"].parent
+        for k, v in enumerate(pred.predictions):
+            row[f"prediction_label_{k}"] = v[0]
+            row[f"confidence__label_{k}"] = v[1]
+        with open(csv_fpath, mode="a", newline="") as csv_file:
+            field_names = list(row.keys())
+            writer = csv.DictWriter(csv_file, fieldnames=field_names)
+            writer.writerow(row)
+
+
 class InferenceSession(ABC):
     """Abstract class for inference sessions"""
 
     @abstractmethod
-    def __init__(self, model: Path, vocab: List):
+    def __init__(self, model: Path, vocab: List):  # pragma: no cover
         """Inference Sessions should init from a model file and vocab"""
         self.model = model
         self.vocab = vocab
 
     @abstractmethod
-    def predict_image(self, image: Path):
+    def predict_image(self, image: Path):  # pragma: no cover
         """Predict a single image"""
         pass
 
     @abstractmethod
-    def predict_batch(self, model: Path, batch: Iterable[Path], bs: int):
+    def predict_batch(
+        self, model: Path, batch: Iterable[Path], bs: int
+    ):  # pragma: no cover
         """Predict a batch"""
         pass
 
@@ -232,7 +361,7 @@ def softmax(x):
 #         return PredictionBatch(prediction_items)
 
 
-class OnnxInferenceSession(InferenceSession):
+class OnnxInferenceSession(InferenceSession):  # pragma: no cover
     """onnx inference session"""
 
     def __init__(self, model: Path, vocab: Path):
@@ -240,22 +369,31 @@ def __init__(self, model: Path, vocab: Path):
         self.model = model
         self.session = rt.InferenceSession(str(model))
 
-        self.vocab = vocab
-        self.vocab_mapping = dict(enumerate(self.vocab))
-
-    def _load_vocab(self, vocab: Path) -> List:
-        with open(vocab, "r") as f:
-            return [item.strip("\n") for item in f.readlines()]
+        self.vocab = models.load_vocab(vocab)
+        self.vocab_mappings = [dict(enumerate(v)) for v in self.vocab]
 
-    def predict_image(self, image: Path):
+    def predict_image(
+        self, image: Path
+    ) -> Union[ImagePredictionItem, MultiLabelImagePredictionItem]:
         """Predict a single image"""
         img = self._load_image(image)
-        raw_result = self.session.run(["output"], {"image": img})
-        pred = self._postprocess(raw_result)
-        arg_max = int(np.array(pred).argmax())
-        predicted_label = self.vocab_mapping[int(arg_max)]
-        confidence = float(np.array(pred).max())
-        return ImagePredictionItem(image, predicted_label, confidence)
+        output_names = [o.name for o in self.session.get_outputs()]
+        raw_result = self.session.run(output_names, {"image": img})
+        if len(self.vocab) < 2:
+            pred = self._postprocess(raw_result)
+            arg_max = int(np.array(pred).argmax())
+            predicted_label = self.vocab_mappings[0][int(arg_max)]
+            confidence = float(np.array(pred).max())
+            return ImagePredictionItem(image, predicted_label, confidence)
+        else:
+            prediction_tuples = []
+            for vocab_map, pred in zip(self.vocab_mappings, raw_result):
+                pred = self._postprocess(pred)
+                arg_max = int(np.array(pred).argmax())
+                predicted_label = vocab_map[int(arg_max)]
+                confidence = float(np.array(pred).max())
+                prediction_tuples.append((predicted_label, confidence))
+        return MultiLabelImagePredictionItem(image, prediction_tuples)
 
     def _preprocess(self, input_data: np.ndarray) -> np.ndarray:
         # converts the input data into the float32 input for onnx
@@ -285,11 +423,16 @@ def _postprocess(self, result: List):
         """process results from onnx session"""
         return softmax(np.array(result)).tolist()
 
-    def predict_batch(self, batch: Iterable[Path], bs: int):
+    def predict_batch(
+        self, batch: Iterable[Path], bs: int
+    ) -> Union[PredictionBatch, MultiPredictionBatch]:
         """predicts a batch of images"""
         prediction_items = [self.predict_image(file) for file in batch]
-        return PredictionBatch(prediction_items)
+        if len(self.vocab) < 2:
+            return PredictionBatch(prediction_items)
+        else:
+            return MultiPredictionBatch(prediction_items)
 
 
 if __name__ == "__main__":
-    app()
+    app()  # pragma: no cover
diff --git a/src/flyswot/models.py b/src/flyswot/models.py
--- a/src/flyswot/models.py
+++ b/src/flyswot/models.py
@@ -11,17 +11,19 @@
 from typing import Dict
 from typing import List
 from typing import Optional
+from typing import Tuple
 from typing import Union
 
 import typer
 import validators  # type: ignore
 from rich.markdown import Markdown
+from toolz import itertoolz
+from toolz import recipes
 
 from flyswot.config import APP_NAME
 from flyswot.config import MODEL_REPO_URL
 from flyswot.console import console
 
-
 app = typer.Typer()
 
 REPO_URL = "https://api.github.com/repos/davanstrien/learn-onnx/releases"
@@ -42,9 +44,9 @@ class GitHubRelease:
 class LocalModel:
     """A local model container"""
 
-    vocab: Optional[Path]
-    modelcard: Optional[Path]
-    model: Optional[Path]
+    vocab: Path
+    modelcard: Path
+    model: Path
 
 
 def get_release_metadata(release: Dict[str, Any]) -> GitHubRelease:
@@ -70,7 +72,9 @@ def _url_callback(url: str) -> Union[str, None]:
         raise typer.BadParameter(f"Please check {url} is a valid url")
 
 
-def get_remote_release_json(url: str, single: bool = False) -> Dict[str, Any]:
+def get_remote_release_json(
+    url: str, single: bool = False
+) -> Dict[str, Any]:  # pragma : no cover
     """Returns all releases found at `url`"""
     with urllib.request.urlopen(url) as response:  # pragma: no cover
         try:
@@ -167,7 +171,6 @@ def _get_latest_model(model_dir: Path) -> Optional[Path]:
 def _get_model_parts(model_dir: Path) -> LocalModel:
     """Returns model path, vocab and metadata for a model"""
     model_files = Path(model_dir / "model").iterdir()
-    vocab, modelcard, model = None, None, None
     for file in model_files:
         if fnmatch.fnmatch(file.name, "vocab.txt"):
             vocab = file
@@ -190,6 +193,11 @@ def _compare_remote_local(local_model: Path) -> bool:  # pragma: no cover
     )
 
 
+def load_model(model_dir: Path):  # pragma: no cover
+    """returns a local model"""
+    return _get_model_parts(model_dir)
+
+
 def ensure_model(
     model_dir: Path, check_latest: bool = False
 ) -> Optional[LocalModel]:  # pragma: no cover
@@ -218,10 +226,21 @@ def ensure_model(
         raise typer.Exit()
 
 
-def load_vocab(vocab: Path) -> List[str]:
-    """loads vocab from `vocab` and returns as list"""
+def is_pipe(c: Tuple) -> bool:
+    """Checks if | in c"""
+    return "|" in c
+
+
+def load_vocab(vocab: Path) -> List[List[str]]:
+    """loads vocab from `vocab` and returns as list contaning lists of vocab"""
     with open(vocab, "r") as f:
-        return [line.strip("\n") for line in f.readlines()]
+        raw_vocab = [line.strip("\n") for line in f.readlines()]
+        return list(
+            map(
+                list,
+                (itertoolz.remove(is_pipe, (recipes.partitionby(is_pipe, raw_vocab)))),
+            )
+        )
 
 
 @app.command()
@@ -241,9 +260,6 @@ def vocab(
                 console.print(Markdown("# Model Vocab"))
                 console.print(vocab)
             return vocab
-    if not model_path:
-        typer.echo(f"No models currently found in {model_path}")
-        raise typer.Exit()
 
 
 if __name__ == "__main__":  # pragma: no cover


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -25,6 +25,8 @@ jobs:
     steps:
       - name: Check out the repository
         uses: actions/checkout@v2.3.4
+        with:
+          lfs: true
 
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2.2.2
diff --git a/tests/test_files/mult/20210629/model/2021-06-29-model.onnx b/tests/test_files/mult/20210629/model/2021-06-29-model.onnx
new file mode 100644
--- /dev/null
+++ b/tests/test_files/mult/20210629/model/2021-06-29-model.onnx
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84c8fde3fd9175e99a3580997129b342a808a67b415eea6106d93a14aa59c081
+size 108756013
diff --git a/tests/test_files/mult/20210629/model/modelcard.md b/tests/test_files/mult/20210629/model/modelcard.md
new file mode 100644
--- /dev/null
+++ b/tests/test_files/mult/20210629/model/modelcard.md
@@ -0,0 +1,19 @@
+# flyswot
+
+## Model description
+
+In progress model for detecting 'fake' flysheets
+
+## Intended uses & limitations
+
+Not currently intended for public consumption...
+
+#### Limitations and bias
+
+Not currently intended for public consumption...
+
+## Training data
+
+TODO
+
+## Eval results
diff --git a/tests/test_files/mult/20210629/model/vocab.txt b/tests/test_files/mult/20210629/model/vocab.txt
new file mode 100644
--- /dev/null
+++ b/tests/test_files/mult/20210629/model/vocab.txt
@@ -0,0 +1,18 @@
+flysheet
+other
+|
+Back flysheet
+Back folder
+Box
+Carrying case or wrapper
+Control shot
+Front flysheet
+Front folder
+Material wrapper
+Original bindings
+Other papers letters cards
+Pages
+Scroll images
+Second foliation sequence
+cover
+edge
diff --git a/tests/test_files/test_vocab.txt b/tests/test_files/test_vocab.txt
new file mode 100644
--- /dev/null
+++ b/tests/test_files/test_vocab.txt
@@ -0,0 +1,18 @@
+flysheet
+other
+|
+Back flysheet
+Back folder
+Box
+Carrying case or wrapper
+Control shot
+Front flysheet
+Front folder
+Material wrapper
+Original bindings
+Other papers letters cards
+Pages
+Scroll images
+Second foliation sequence
+cover
+edge
diff --git a/tests/test_inference.py b/tests/test_inference.py
--- a/tests/test_inference.py
+++ b/tests/test_inference.py
@@ -1,20 +1,27 @@
 import csv
+import inspect
 import os
 import pathlib
+from collections import defaultdict
+from pathlib import Path
 from typing import Any
 
 import pytest
+import rich
 import typer
 from hypothesis import given
+from hypothesis import strategies
 from hypothesis import strategies as st
 from hypothesis.core import given
+from rich.table import Column
+from toolz import itertoolz
 
 from flyswot import inference
 
 # flake8: noqa
 
 
-def test__image_predicton_item_raises_error_when_too_few_params():
+def test__image_prediction_item_raises_error_when_too_few_params():
     """It raises Typerror when passed too few items"""
     with pytest.raises(TypeError):
         item = inference.ImagePredictionItem("A")
@@ -49,6 +56,20 @@ def test_image_prediction_item(confidence, label, imfile: Any):
     assert type(item.confidence) == float
 
 
+@given(confidence=st.floats(max_value=100.0), label=st.text(min_size=1))
+def test_multi_image_prediction_item(confidence, label, imfile: Any):
+    item = inference.MultiLabelImagePredictionItem(
+        imfile, [(label, confidence), (label, confidence)]
+    )
+    assert item.path == imfile
+    assert type(item.predictions) == list
+    assert type(item.predictions[0]) == tuple
+    assert item.predictions[0][0] == label
+    assert item.predictions[1][0] == label
+    assert item.predictions[0][1] == confidence
+    assert item.predictions[1][1] == confidence
+
+
 @given(confidence=st.floats(max_value=100.0), label=st.text(min_size=1))
 def test_prediction_batch(confidence: float, label: str, imfile: Any):
     item = inference.ImagePredictionItem(imfile, label, confidence)
@@ -59,6 +80,26 @@ def test_prediction_batch(confidence: float, label: str, imfile: Any):
     assert hasattr(batch.batch_labels, "__next__")
 
 
+@given(confidence=st.floats(min_value=0.0, max_value=100.0), label=st.text(min_size=1))
+def test_multi_prediction_batch(confidence: float, label: str, imfile: Any):
+    item = inference.MultiLabelImagePredictionItem(
+        imfile, [(label, confidence), (label, confidence)]
+    )
+    item2 = inference.MultiLabelImagePredictionItem(
+        imfile, [(label, confidence), (label, confidence)]
+    )
+    batch = inference.MultiPredictionBatch([item, item2])
+    assert batch.batch
+    assert type(batch.batch) == list
+    assert type(batch.batch[0]) == inference.MultiLabelImagePredictionItem
+    assert batch.batch_labels
+    assert len(list(batch.batch_labels)) == 2
+    assert hasattr(batch.batch_labels, "__next__")
+    for labels in batch.batch_labels:
+        for label in labels:
+            assert label == label
+
+
 FIXTURE_DIR = os.path.join(
     os.path.dirname(os.path.realpath(__file__)),
     "test_files",
@@ -73,7 +114,56 @@ def test_prediction_batch(confidence: float, label: str, imfile: Any):
 )
 def test_predict_directory(datafiles, tmp_path) -> None:
     inference.predict_directory(
-        datafiles, tmp_path, pattern="fse", bs=1, image_format=".jpg"
+        datafiles,
+        tmp_path,
+        pattern="fse",
+        bs=1,
+        image_format=".jpg",
+        model_name="latest",
+    )
+    csv_file = list(tmp_path.rglob("*.csv"))
+    assert csv_file
+    with open(csv_file[0], newline="") as csvfile:
+        reader = csv.DictReader(csvfile)
+        for row in reader:
+            assert row["path"]
+            assert row["directory"]
+        columns = defaultdict(list)
+    with open(csv_file[0], newline="") as csvfile:
+        reader = csv.DictReader(csvfile)
+
+        for row in reader:
+            for (k, v) in row.items():
+                columns[k].append(v)
+        assert any("prediction" in k for k in columns)
+        labels = [columns[k] for k in columns if "prediction" in k]
+        confidences = [columns[k] for k in columns if "confidence" in k]
+        # check all labels are strings
+        assert all(map(lambda x: isinstance(x, str), (itertoolz.concat(labels))))
+        # check all confidences can be cast to float
+        assert all(
+            map(
+                lambda x: isinstance(x, float),
+                map(lambda x: float(x), (itertoolz.concat(confidences))),
+            )
+        )
+
+
+@pytest.mark.datafiles(
+    os.path.join(
+        FIXTURE_DIR,
+        "fly_fse.jpg",
+    )
+)
+def test_predict_directory_local_mult(datafiles, tmp_path) -> None:
+    inference.predict_directory(
+        datafiles,
+        tmp_path,
+        pattern="fse",
+        bs=1,
+        image_format=".jpg",
+        model_name="20210629",
+        model_path="tests/test_files/mult/",
     )
     csv_file = list(tmp_path.rglob("*.csv"))
     assert csv_file
@@ -82,13 +172,116 @@ def test_predict_directory(datafiles, tmp_path) -> None:
         for row in reader:
             assert row["path"]
             assert row["directory"]
-            assert row["predicted_label"]
-            assert row["confidence"]
-            assert type(float(row["confidence"])) == float
+        columns = defaultdict(list)
+    with open(csv_file[0], newline="") as csvfile:
+        reader = csv.DictReader(csvfile)
+
+        for row in reader:
+            for (k, v) in row.items():
+                columns[k].append(v)
+        assert any("prediction" in k for k in columns)
+        labels = [columns[k] for k in columns if "prediction" in k]
+        confidences = [columns[k] for k in columns if "confidence" in k]
+        # check all labels are strings
+        assert all(map(lambda x: isinstance(x, str), (itertoolz.concat(labels))))
+        # check all confidences can be cast to float
+        assert all(
+            map(
+                lambda x: isinstance(x, float),
+                map(lambda x: float(x), (itertoolz.concat(confidences))),
+            )
+        )
+
+
+def test_csv_header():
+    with pytest.raises(NotImplementedError):
+        inference.create_csv_header("string", Path("."))
+
+
+def test_csv_header_multi(tmp_path):
+    predicton = inference.MultiLabelImagePredictionItem(Path("."), [("label", 0.8)])
+    batch = inference.MultiPredictionBatch([predicton, predicton])
+    csv_fname = tmp_path / "test.csv"
+    inference.create_csv_header(batch, csv_fname)
+    with open(csv_fname, "r") as f:
+        reader = csv.DictReader(f)
+        list_of_column_names = [header for header in reader.fieldnames]
+    assert "path" in list_of_column_names
+    assert "directory" in list_of_column_names
+
+
+def test_csv_header_single(tmp_path):
+    predicton = inference.ImagePredictionItem(Path("."), "label", 0.6)
+    batch = inference.PredictionBatch([predicton])
+    csv_fname = tmp_path / "test.csv"
+    inference.create_csv_header(batch, csv_fname)
+    with open(csv_fname, "r") as f:
+        reader = csv.DictReader(f)
+        list_of_column_names = [header for header in reader.fieldnames]
+    assert "path" in list_of_column_names
+    assert "directory" in list_of_column_names
+    assert "confidence" in list_of_column_names
+
+
+def test_csv_batch():
+    with pytest.raises(NotImplementedError):
+        inference.write_batch_preds_to_csv("string", Path("."))
 
 
-def test_predict_empty_directory(tmp_path) -> None:
-    test_dir = tmp_path / "test"
-    test_dir.mkdir()
+def test_csv_batch_single(tmp_path):
+    predicton = inference.ImagePredictionItem(Path("."), "label", 0.6)
+    batch = inference.PredictionBatch([predicton])
+    csv_fname = tmp_path / "test.csv"
+    inference.write_batch_preds_to_csv(batch, csv_fname)
+    with open(csv_fname, "r") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            assert "label" in row
+
+
+@given(strategies.lists(st.text(min_size=2), min_size=2), st.text())
+def test_print_table(labels, title):
+    table = inference.print_table(labels, title, print=False)
+    assert isinstance(table, rich.table.Table)
+    assert table.title == title
+    unique = itertoolz.count(itertoolz.unique(labels))
+    assert table.row_count == unique + 1
+    assert all(
+        label in getattr(itertoolz.first(table.columns), "_cells") for label in labels
+    )
+    table = inference.print_table(labels, title, print=True)
+
+
+@given(strategies.lists(st.text(), min_size=1))
+def test_check_files(l):
+    inference.check_files(l, "fse", Path("."))
+
+
+def test_check_files_emopty():
     with pytest.raises(typer.Exit):
-        inference.predict_directory(test_dir)
+        inference.check_files([], "fse", Path("."))
+
+
+def test_make_layout():
+    layout = inference.make_layout()
+    assert layout
+    assert isinstance(layout, rich.layout.Layout)
+
+
+# def test_predict_empty_directory(tmp_path) -> None:
+#     test_dir = tmp_path / "test"
+#     test_dir.mkdir()
+#     with pytest.raises(typer.Exit):
+#         inference.predict_directory(test_dir)
+
+
+def test_onnx_session_attributes():
+    session = inference.OnnxInferenceSession(
+        Path("tests/test_files/mult/20210629/model/2021-06-29-model.onnx"),
+        Path("tests/test_files/mult/20210629/model/vocab.txt"),
+    )
+    assert session
+    assert inspect.ismethod(session.predict_image)
+    assert inspect.ismethod(session.predict_batch)
+    assert inspect.ismethod(session._postprocess)
+    assert inspect.ismethod(session._load_image)
diff --git a/tests/test_models.py b/tests/test_models.py
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -15,6 +15,8 @@
 from flyswot import models
 from flyswot.models import app
 from flyswot.models import ensure_model_dir
+from flyswot.models import load_vocab
+from flyswot.models import vocab
 
 runner = CliRunner()
 
@@ -198,5 +200,30 @@ def test_ensure_model() -> None:
     assert model_parts
 
 
+def test_load_vocab() -> None:
+    vocab = models.load_vocab(Path("tests/test_files/test_vocab.txt"))
+    assert vocab
+    assert type(vocab) == list
+    assert type(vocab[0]) == list
+
+
+def test_vocab_raises() -> None:
+    with pytest.raises(NotImplementedError):
+        models.vocab(model="not_latest")
+
+
+def test_vocab() -> None:
+    vocab = models.vocab(model="latest", show=False)
+    assert vocab
+    assert isinstance(vocab, list)
+
+
+def test_vocab_print(capsys) -> None:
+    models.vocab(model="latest", show=True)
+    captured = capsys.readouterr()
+    assert len(captured.out) > 2
+    assert "Model Vocab" in captured.out
+
+
 def test_app(tmp_path: Any) -> None:
     pass