We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
update default for ensure_model to get latest model 


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/flyswot/inference.py`
Content:
```
1 """Inference functionality"""
2 import csv
3 import mimetypes
4 import time
5 from abc import ABC
6 from abc import abstractmethod
7 from dataclasses import asdict
8 from dataclasses import dataclass
9 from datetime import datetime
10 from datetime import timedelta
11 from pathlib import Path
12 from typing import Iterable
13 from typing import Iterator
14 from typing import List
15 from typing import Union
16 
17 import numpy as np
18 import onnxruntime as rt  # type: ignore
19 import typer
20 from PIL import Image  # type: ignore
21 from rich.table import Table
22 from toolz import itertoolz  # type: ignore
23 
24 from flyswot import core
25 from flyswot import models
26 from flyswot.console import console
27 
28 app = typer.Typer()
29 
30 
31 @dataclass()
32 class ImagePredictionItem:
33     """Prediction for an image.
34 
35     Attributes:
36         path: The Path to the image
37         predicted_label: The predicted label i.e. the argmax value for the prediction tensor
38         condidence: The confidence for `predicted_label` i.e. the max value for prediction tensor
39     """
40 
41     path: Path
42     predicted_label: str
43     confidence: float
44 
45     def __post_init__(self) -> Union[Path, None]:
46         """attempt to get absolute path"""
47         try:
48             self.path: Path = self.path.absolute()
49         except AttributeError:
50             pass
51 
52 
53 @dataclass
54 class PredictionBatch:
55     """Container for ImagePredictionItems"""
56 
57     batch: List[ImagePredictionItem]
58 
59     def __post_init__(self):
60         """Returns a list of all predicted labels in batch"""
61         self.batch_labels: Iterator[str] = (item.predicted_label for item in self.batch)
62 
63 
64 image_extensions = {k for k, v in mimetypes.types_map.items() if v.startswith("image/")}
65 
66 
67 @app.command()
68 def predict_image(
69     image: Path = typer.Argument(..., readable=True, resolve_path=True)
70 ) -> None:
71     """Predict a single image"""
72     pass  # pragma: no cover
73 
74 
75 @app.command(name="directory")
76 def predict_directory(
77     directory: Path = typer.Argument(
78         ...,
79         readable=True,
80         resolve_path=True,
81         help="Directory to start searching for images from",
82     ),
83     csv_save_dir: Path = typer.Argument(
84         ...,
85         writable=True,
86         resolve_path=True,
87         help="Directory used to store the csv report",
88     ),
89     pattern: str = typer.Option("fse", help="Pattern used to filter image filenames"),
90     bs: int = typer.Option(16, help="Batch Size"),
91     image_format: str = typer.Option(
92         ".tif",
93         help="Image format for flyswot to use for predictions",
94     ),
95 ):
96     """Predicts against all images stored under DIRECTORY which match PATTERN in the filename.
97 
98     By default searches for filenames containing 'fse'.
99 
100     Creates a CSV report saved to `csv_save_dir`
101     """
102     start_time = time.perf_counter()
103     model_dir = models.ensure_model_dir()
104     # TODO add load learner function that can be passed a model name
105     model_parts = models.ensure_model(model_dir)
106     model = model_parts.model
107     vocab = models.load_vocab(model_parts.vocab)
108     onnxinference = OnnxInferenceSession(model, vocab)
109     files = core.get_image_files_from_pattern(directory, pattern, image_format)
110     files = list(files)
111     typer.echo(f"Found {len(files)} files matching {pattern} in {directory}")
112     csv_fname = create_csv_fname(csv_save_dir)
113     create_csv_header(csv_fname)
114     with typer.progressbar(length=len(files)) as progress:
115         all_preds = []
116         predictions = []
117         for batch in itertoolz.partition_all(bs, files):
118             batch_predictions = onnxinference.predict_batch(batch, bs)
119             all_preds.append(batch_predictions.batch_labels)
120             predictions.append(batch_predictions)
121             progress.update(len(batch))
122             write_batch_preds_to_csv(csv_fname, batch_predictions)
123         all_preds = list(itertoolz.concat(all_preds))
124     typer.echo(f"CSV report stored in {csv_fname}")
125     delta = timedelta(seconds=time.perf_counter() - start_time)
126     typer.echo(f"Time taken to run:  {str(delta)}")
127     print_table(all_preds)
128 
129 
130 def print_table(decoded) -> None:
131     """Prints table summary of predicted labels"""
132     table = Table(show_header=True, title="Prediction summary")
133     table.add_column(
134         "Class",
135     )
136     table.add_column("Count")
137     table.add_column("Percentage")
138     total = len(decoded)
139     frequencies = itertoolz.frequencies(decoded)
140     for is_last_element, var in core.signal_last(frequencies.items()):
141         key, value = var
142         count = value
143         percentage = round((count / total) * 100, 2)
144         if is_last_element:
145             table.add_row(key, str(count), f"{percentage}", end_section=True)
146             table.add_row("Total", str(total), "")
147         else:
148             table.add_row(key, str(count), f"{percentage}")
149     console.print(table)
150 
151 
152 def create_csv_fname(csv_directory: Path) -> Path:
153     """Creates a csv filename"""
154     date_now = datetime.now()
155     date_now = date_now.strftime("%Y_%m_%d_%H_%M")
156     fname = Path(date_now + ".csv")
157     return Path(csv_directory / fname)
158 
159 
160 def create_csv_header(csv_path: Path) -> None:
161     """Creates a header for csv `csv_path`"""
162     with open(csv_path, mode="w", newline="") as csv_file:
163         field_names = ["path", "directory", "predicted_label", "confidence"]
164         writer = csv.DictWriter(csv_file, fieldnames=field_names)
165         writer.writeheader()
166 
167 
168 def write_batch_preds_to_csv(csv_fpath: Path, predictions: PredictionBatch) -> None:
169     """Appends `predictions` batch to `csv_path`"""
170     with open(csv_fpath, mode="a", newline="") as csv_file:
171         field_names = ["path", "directory", "predicted_label", "confidence"]
172         writer = csv.DictWriter(csv_file, fieldnames=field_names)
173         for pred in predictions.batch:
174             row = asdict(pred)
175             row["directory"] = pred.path.parent
176             writer.writerow(row)
177 
178 
179 class InferenceSession(ABC):
180     """Abstract class for inference sessions"""
181 
182     @abstractmethod
183     def __init__(self, model: Path, vocab: List):
184         """Inference Sessions should init from a model file and vocab"""
185         self.model = model
186         self.vocab = vocab
187 
188     @abstractmethod
189     def predict_image(self, image: Path):
190         """Predict a single image"""
191         pass
192 
193     @abstractmethod
194     def predict_batch(self, model: Path, batch: Iterable[Path], bs: int):
195         """Predict a batch"""
196         pass
197 
198 
199 def softmax(x):
200     """return softmax of `x`"""
201     x = x.reshape(-1)
202     e_x = np.exp(x - np.max(x))
203     return e_x / e_x.sum(axis=0)
204 
205 
206 # class FastaiInferenceModel(InferenceSession):
207 #     def __init__(self, model):
208 #         self.model = model
209 #         self.learn = load_learner(model)
210 
211 #     def predict_image(self, image: Path) -> Any:
212 #         return self.learn.predict(image)
213 
214 #     def predict_batch(self, batch: Iterable[Path], bs: int) -> PredictionBatch:
215 #         test_dl = self.learn.dls.test_dl(batch, bs=bs)
216 #         vocab = dict(enumerate(self.learn.dls.vocab))
217 #         with self.learn.no_bar():
218 #             fastai_preds: Any = self.learn.get_preds(dl=test_dl, with_decoded=True)
219 #             prediction_tensors: Iterable[Any] = fastai_preds[0]
220 #             prediction_items = []
221 #             for file, pred in zip(batch, prediction_tensors):
222 #                 arg_max = int(np.array(pred).argmax())
223 #                 predicted_label = vocab[int(arg_max)]
224 #                 confidence = float(np.array(pred).max())
225 #                 prediction_items.append(
226 #                     ImagePredictionItem(file, predicted_label, confidence)
227 #                 )
228 #         return PredictionBatch(prediction_items)
229 
230 
231 class OnnxInferenceSession(InferenceSession):
232     """onnx inference session"""
233 
234     def __init__(self, model: Path, vocab: Path):
235         """Create onnx session"""
236         self.model = model
237         self.session = rt.InferenceSession(str(model))
238 
239         self.vocab = vocab
240         self.vocab_mapping = dict(enumerate(self.vocab))
241 
242     def _load_vocab(self, vocab: Path) -> List:
243         with open(vocab, "r") as f:
244             return [item.strip("\n") for item in f.readlines()]
245 
246     def predict_image(self, image: Path):
247         """Predict a single image"""
248         img = self._load_image(image)
249         raw_result = self.session.run(["output"], {"image": img})
250         pred = self._postprocess(raw_result)
251         arg_max = int(np.array(pred).argmax())
252         predicted_label = self.vocab_mapping[int(arg_max)]
253         confidence = float(np.array(pred).max())
254         return ImagePredictionItem(image, predicted_label, confidence)
255 
256     def _preprocess(self, input_data: np.ndarray) -> np.ndarray:
257         # converts the input data into the float32 input for onnx
258         img_data = input_data.astype("float32")
259 
260         # normalize
261         mean_vec = np.array([0.485, 0.456, 0.406])
262         stddev_vec = np.array([0.229, 0.224, 0.225])
263         norm_img_data = np.zeros(img_data.shape).astype("float32")
264         for i in range(img_data.shape[0]):
265             norm_img_data[i, :, :] = (
266                 img_data[i, :, :] / 255 - mean_vec[i]
267             ) / stddev_vec[i]
268 
269         # add batch channel
270         norm_img_data = norm_img_data.reshape(1, 3, 512, 512).astype("float32")
271         return norm_img_data
272 
273     def _load_image(self, file: Path) -> np.ndarray:
274         """loads image and carries out preprocessing for inference"""
275         image = Image.open(file, mode="r")
276         image = image.resize((512, 512), Image.BILINEAR)
277         image_data = np.array(image).transpose(2, 0, 1)
278         return self._preprocess(image_data)
279 
280     def _postprocess(self, result: List):
281         """process results from onnx session"""
282         return softmax(np.array(result)).tolist()
283 
284     def predict_batch(self, batch: Iterable[Path], bs: int):
285         """predicts a batch of images"""
286         prediction_items = [self.predict_image(file) for file in batch]
287         return PredictionBatch(prediction_items)
288 
289 
290 if __name__ == "__main__":
291     app()
292
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/flyswot/inference.py b/src/flyswot/inference.py
--- a/src/flyswot/inference.py
+++ b/src/flyswot/inference.py
@@ -19,7 +19,7 @@
 import typer
 from PIL import Image  # type: ignore
 from rich.table import Table
-from toolz import itertoolz  # type: ignore
+from toolz import itertoolz
 
 from flyswot import core
 from flyswot import models
@@ -89,9 +89,9 @@
     pattern: str = typer.Option("fse", help="Pattern used to filter image filenames"),
     bs: int = typer.Option(16, help="Batch Size"),
     image_format: str = typer.Option(
-        ".tif",
-        help="Image format for flyswot to use for predictions",
+        ".tif", help="Image format for flyswot to use for predictions"
     ),
+    check_latest: bool = typer.Option(True, help="Use latest available model"),
 ):
     """Predicts against all images stored under DIRECTORY which match PATTERN in the filename.
 
@@ -102,7 +102,7 @@
     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)
+    model_parts = models.ensure_model(model_dir, check_latest)
     model = model_parts.model
     vocab = models.load_vocab(model_parts.vocab)
     onnxinference = OnnxInferenceSession(model, vocab)

{"golden_diff": "diff --git a/src/flyswot/inference.py b/src/flyswot/inference.py\n--- a/src/flyswot/inference.py\n+++ b/src/flyswot/inference.py\n@@ -19,7 +19,7 @@\n import typer\n from PIL import Image  # type: ignore\n from rich.table import Table\n-from toolz import itertoolz  # type: ignore\n+from toolz import itertoolz\n \n from flyswot import core\n from flyswot import models\n@@ -89,9 +89,9 @@\n     pattern: str = typer.Option(\"fse\", help=\"Pattern used to filter image filenames\"),\n     bs: int = typer.Option(16, help=\"Batch Size\"),\n     image_format: str = typer.Option(\n-        \".tif\",\n-        help=\"Image format for flyswot to use for predictions\",\n+        \".tif\", help=\"Image format for flyswot to use for predictions\"\n     ),\n+    check_latest: bool = typer.Option(True, help=\"Use latest available model\"),\n ):\n     \"\"\"Predicts against all images stored under DIRECTORY which match PATTERN in the filename.\n \n@@ -102,7 +102,7 @@\n     start_time = time.perf_counter()\n     model_dir = models.ensure_model_dir()\n     # TODO add load learner function that can be passed a model name\n-    model_parts = models.ensure_model(model_dir)\n+    model_parts = models.ensure_model(model_dir, check_latest)\n     model = model_parts.model\n     vocab = models.load_vocab(model_parts.vocab)\n     onnxinference = OnnxInferenceSession(model, vocab)\n", "issue": "update default for ensure_model to get latest model \n\n", "before_files": [{"content": "\"\"\"Inference functionality\"\"\"\nimport csv\nimport mimetypes\nimport time\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom dataclasses import asdict\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom pathlib import Path\nfrom typing import Iterable\nfrom typing import Iterator\nfrom typing import List\nfrom typing import Union\n\nimport numpy as np\nimport onnxruntime as rt  # type: ignore\nimport typer\nfrom PIL import Image  # type: ignore\nfrom rich.table import Table\nfrom toolz import itertoolz  # type: ignore\n\nfrom flyswot import core\nfrom flyswot import models\nfrom flyswot.console import console\n\napp = typer.Typer()\n\n\n@dataclass()\nclass ImagePredictionItem:\n    \"\"\"Prediction for an image.\n\n    Attributes:\n        path: The Path to the image\n        predicted_label: The predicted label i.e. the argmax value for the prediction tensor\n        condidence: The confidence for `predicted_label` i.e. the max value for prediction tensor\n    \"\"\"\n\n    path: Path\n    predicted_label: str\n    confidence: float\n\n    def __post_init__(self) -> Union[Path, None]:\n        \"\"\"attempt to get absolute path\"\"\"\n        try:\n            self.path: Path = self.path.absolute()\n        except AttributeError:\n            pass\n\n\n@dataclass\nclass PredictionBatch:\n    \"\"\"Container for ImagePredictionItems\"\"\"\n\n    batch: List[ImagePredictionItem]\n\n    def __post_init__(self):\n        \"\"\"Returns a list of all predicted labels in batch\"\"\"\n        self.batch_labels: Iterator[str] = (item.predicted_label for item in self.batch)\n\n\nimage_extensions = {k for k, v in mimetypes.types_map.items() if v.startswith(\"image/\")}\n\n\n@app.command()\ndef predict_image(\n    image: Path = typer.Argument(..., readable=True, resolve_path=True)\n) -> None:\n    \"\"\"Predict a single image\"\"\"\n    pass  # pragma: no cover\n\n\n@app.command(name=\"directory\")\ndef predict_directory(\n    directory: Path = typer.Argument(\n        ...,\n        readable=True,\n        resolve_path=True,\n        help=\"Directory to start searching for images from\",\n    ),\n    csv_save_dir: Path = typer.Argument(\n        ...,\n        writable=True,\n        resolve_path=True,\n        help=\"Directory used to store the csv report\",\n    ),\n    pattern: str = typer.Option(\"fse\", help=\"Pattern used to filter image filenames\"),\n    bs: int = typer.Option(16, help=\"Batch Size\"),\n    image_format: str = typer.Option(\n        \".tif\",\n        help=\"Image format for flyswot to use for predictions\",\n    ),\n):\n    \"\"\"Predicts against all images stored under DIRECTORY which match PATTERN in the filename.\n\n    By default searches for filenames containing 'fse'.\n\n    Creates a CSV report saved to `csv_save_dir`\n    \"\"\"\n    start_time = time.perf_counter()\n    model_dir = models.ensure_model_dir()\n    # TODO add load learner function that can be passed a model name\n    model_parts = models.ensure_model(model_dir)\n    model = model_parts.model\n    vocab = models.load_vocab(model_parts.vocab)\n    onnxinference = OnnxInferenceSession(model, vocab)\n    files = core.get_image_files_from_pattern(directory, pattern, image_format)\n    files = list(files)\n    typer.echo(f\"Found {len(files)} files matching {pattern} in {directory}\")\n    csv_fname = create_csv_fname(csv_save_dir)\n    create_csv_header(csv_fname)\n    with typer.progressbar(length=len(files)) as progress:\n        all_preds = []\n        predictions = []\n        for batch in itertoolz.partition_all(bs, files):\n            batch_predictions = onnxinference.predict_batch(batch, bs)\n            all_preds.append(batch_predictions.batch_labels)\n            predictions.append(batch_predictions)\n            progress.update(len(batch))\n            write_batch_preds_to_csv(csv_fname, batch_predictions)\n        all_preds = list(itertoolz.concat(all_preds))\n    typer.echo(f\"CSV report stored in {csv_fname}\")\n    delta = timedelta(seconds=time.perf_counter() - start_time)\n    typer.echo(f\"Time taken to run:  {str(delta)}\")\n    print_table(all_preds)\n\n\ndef print_table(decoded) -> None:\n    \"\"\"Prints table summary of predicted labels\"\"\"\n    table = Table(show_header=True, title=\"Prediction summary\")\n    table.add_column(\n        \"Class\",\n    )\n    table.add_column(\"Count\")\n    table.add_column(\"Percentage\")\n    total = len(decoded)\n    frequencies = itertoolz.frequencies(decoded)\n    for is_last_element, var in core.signal_last(frequencies.items()):\n        key, value = var\n        count = value\n        percentage = round((count / total) * 100, 2)\n        if is_last_element:\n            table.add_row(key, str(count), f\"{percentage}\", end_section=True)\n            table.add_row(\"Total\", str(total), \"\")\n        else:\n            table.add_row(key, str(count), f\"{percentage}\")\n    console.print(table)\n\n\ndef create_csv_fname(csv_directory: Path) -> Path:\n    \"\"\"Creates a csv filename\"\"\"\n    date_now = datetime.now()\n    date_now = date_now.strftime(\"%Y_%m_%d_%H_%M\")\n    fname = Path(date_now + \".csv\")\n    return Path(csv_directory / fname)\n\n\ndef create_csv_header(csv_path: Path) -> None:\n    \"\"\"Creates a header for csv `csv_path`\"\"\"\n    with open(csv_path, mode=\"w\", newline=\"\") as csv_file:\n        field_names = [\"path\", \"directory\", \"predicted_label\", \"confidence\"]\n        writer = csv.DictWriter(csv_file, fieldnames=field_names)\n        writer.writeheader()\n\n\ndef write_batch_preds_to_csv(csv_fpath: Path, predictions: PredictionBatch) -> None:\n    \"\"\"Appends `predictions` batch to `csv_path`\"\"\"\n    with open(csv_fpath, mode=\"a\", newline=\"\") as csv_file:\n        field_names = [\"path\", \"directory\", \"predicted_label\", \"confidence\"]\n        writer = csv.DictWriter(csv_file, fieldnames=field_names)\n        for pred in predictions.batch:\n            row = asdict(pred)\n            row[\"directory\"] = pred.path.parent\n            writer.writerow(row)\n\n\nclass InferenceSession(ABC):\n    \"\"\"Abstract class for inference sessions\"\"\"\n\n    @abstractmethod\n    def __init__(self, model: Path, vocab: List):\n        \"\"\"Inference Sessions should init from a model file and vocab\"\"\"\n        self.model = model\n        self.vocab = vocab\n\n    @abstractmethod\n    def predict_image(self, image: Path):\n        \"\"\"Predict a single image\"\"\"\n        pass\n\n    @abstractmethod\n    def predict_batch(self, model: Path, batch: Iterable[Path], bs: int):\n        \"\"\"Predict a batch\"\"\"\n        pass\n\n\ndef softmax(x):\n    \"\"\"return softmax of `x`\"\"\"\n    x = x.reshape(-1)\n    e_x = np.exp(x - np.max(x))\n    return e_x / e_x.sum(axis=0)\n\n\n# class FastaiInferenceModel(InferenceSession):\n#     def __init__(self, model):\n#         self.model = model\n#         self.learn = load_learner(model)\n\n#     def predict_image(self, image: Path) -> Any:\n#         return self.learn.predict(image)\n\n#     def predict_batch(self, batch: Iterable[Path], bs: int) -> PredictionBatch:\n#         test_dl = self.learn.dls.test_dl(batch, bs=bs)\n#         vocab = dict(enumerate(self.learn.dls.vocab))\n#         with self.learn.no_bar():\n#             fastai_preds: Any = self.learn.get_preds(dl=test_dl, with_decoded=True)\n#             prediction_tensors: Iterable[Any] = fastai_preds[0]\n#             prediction_items = []\n#             for file, pred in zip(batch, prediction_tensors):\n#                 arg_max = int(np.array(pred).argmax())\n#                 predicted_label = vocab[int(arg_max)]\n#                 confidence = float(np.array(pred).max())\n#                 prediction_items.append(\n#                     ImagePredictionItem(file, predicted_label, confidence)\n#                 )\n#         return PredictionBatch(prediction_items)\n\n\nclass OnnxInferenceSession(InferenceSession):\n    \"\"\"onnx inference session\"\"\"\n\n    def __init__(self, model: Path, vocab: Path):\n        \"\"\"Create onnx session\"\"\"\n        self.model = model\n        self.session = rt.InferenceSession(str(model))\n\n        self.vocab = vocab\n        self.vocab_mapping = dict(enumerate(self.vocab))\n\n    def _load_vocab(self, vocab: Path) -> List:\n        with open(vocab, \"r\") as f:\n            return [item.strip(\"\\n\") for item in f.readlines()]\n\n    def predict_image(self, image: Path):\n        \"\"\"Predict a single image\"\"\"\n        img = self._load_image(image)\n        raw_result = self.session.run([\"output\"], {\"image\": img})\n        pred = self._postprocess(raw_result)\n        arg_max = int(np.array(pred).argmax())\n        predicted_label = self.vocab_mapping[int(arg_max)]\n        confidence = float(np.array(pred).max())\n        return ImagePredictionItem(image, predicted_label, confidence)\n\n    def _preprocess(self, input_data: np.ndarray) -> np.ndarray:\n        # converts the input data into the float32 input for onnx\n        img_data = input_data.astype(\"float32\")\n\n        # normalize\n        mean_vec = np.array([0.485, 0.456, 0.406])\n        stddev_vec = np.array([0.229, 0.224, 0.225])\n        norm_img_data = np.zeros(img_data.shape).astype(\"float32\")\n        for i in range(img_data.shape[0]):\n            norm_img_data[i, :, :] = (\n                img_data[i, :, :] / 255 - mean_vec[i]\n            ) / stddev_vec[i]\n\n        # add batch channel\n        norm_img_data = norm_img_data.reshape(1, 3, 512, 512).astype(\"float32\")\n        return norm_img_data\n\n    def _load_image(self, file: Path) -> np.ndarray:\n        \"\"\"loads image and carries out preprocessing for inference\"\"\"\n        image = Image.open(file, mode=\"r\")\n        image = image.resize((512, 512), Image.BILINEAR)\n        image_data = np.array(image).transpose(2, 0, 1)\n        return self._preprocess(image_data)\n\n    def _postprocess(self, result: List):\n        \"\"\"process results from onnx session\"\"\"\n        return softmax(np.array(result)).tolist()\n\n    def predict_batch(self, batch: Iterable[Path], bs: int):\n        \"\"\"predicts a batch of images\"\"\"\n        prediction_items = [self.predict_image(file) for file in batch]\n        return PredictionBatch(prediction_items)\n\n\nif __name__ == \"__main__\":\n    app()\n", "path": "src/flyswot/inference.py"}], "after_files": [{"content": "\"\"\"Inference functionality\"\"\"\nimport csv\nimport mimetypes\nimport time\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom dataclasses import asdict\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom datetime import timedelta\nfrom pathlib import Path\nfrom typing import Iterable\nfrom typing import Iterator\nfrom typing import List\nfrom typing import Union\n\nimport numpy as np\nimport onnxruntime as rt  # type: ignore\nimport typer\nfrom PIL import Image  # type: ignore\nfrom rich.table import Table\nfrom toolz import itertoolz\n\nfrom flyswot import core\nfrom flyswot import models\nfrom flyswot.console import console\n\napp = typer.Typer()\n\n\n@dataclass()\nclass ImagePredictionItem:\n    \"\"\"Prediction for an image.\n\n    Attributes:\n        path: The Path to the image\n        predicted_label: The predicted label i.e. the argmax value for the prediction tensor\n        condidence: The confidence for `predicted_label` i.e. the max value for prediction tensor\n    \"\"\"\n\n    path: Path\n    predicted_label: str\n    confidence: float\n\n    def __post_init__(self) -> Union[Path, None]:\n        \"\"\"attempt to get absolute path\"\"\"\n        try:\n            self.path: Path = self.path.absolute()\n        except AttributeError:\n            pass\n\n\n@dataclass\nclass PredictionBatch:\n    \"\"\"Container for ImagePredictionItems\"\"\"\n\n    batch: List[ImagePredictionItem]\n\n    def __post_init__(self):\n        \"\"\"Returns a list of all predicted labels in batch\"\"\"\n        self.batch_labels: Iterator[str] = (item.predicted_label for item in self.batch)\n\n\nimage_extensions = {k for k, v in mimetypes.types_map.items() if v.startswith(\"image/\")}\n\n\n@app.command()\ndef predict_image(\n    image: Path = typer.Argument(..., readable=True, resolve_path=True)\n) -> None:\n    \"\"\"Predict a single image\"\"\"\n    pass  # pragma: no cover\n\n\n@app.command(name=\"directory\")\ndef predict_directory(\n    directory: Path = typer.Argument(\n        ...,\n        readable=True,\n        resolve_path=True,\n        help=\"Directory to start searching for images from\",\n    ),\n    csv_save_dir: Path = typer.Argument(\n        ...,\n        writable=True,\n        resolve_path=True,\n        help=\"Directory used to store the csv report\",\n    ),\n    pattern: str = typer.Option(\"fse\", help=\"Pattern used to filter image filenames\"),\n    bs: int = typer.Option(16, help=\"Batch Size\"),\n    image_format: str = typer.Option(\n        \".tif\", help=\"Image format for flyswot to use for predictions\"\n    ),\n    check_latest: bool = typer.Option(True, help=\"Use latest available model\"),\n):\n    \"\"\"Predicts against all images stored under DIRECTORY which match PATTERN in the filename.\n\n    By default searches for filenames containing 'fse'.\n\n    Creates a CSV report saved to `csv_save_dir`\n    \"\"\"\n    start_time = time.perf_counter()\n    model_dir = models.ensure_model_dir()\n    # TODO add load learner function that can be passed a model name\n    model_parts = models.ensure_model(model_dir, check_latest)\n    model = model_parts.model\n    vocab = models.load_vocab(model_parts.vocab)\n    onnxinference = OnnxInferenceSession(model, vocab)\n    files = core.get_image_files_from_pattern(directory, pattern, image_format)\n    files = list(files)\n    typer.echo(f\"Found {len(files)} files matching {pattern} in {directory}\")\n    csv_fname = create_csv_fname(csv_save_dir)\n    create_csv_header(csv_fname)\n    with typer.progressbar(length=len(files)) as progress:\n        all_preds = []\n        predictions = []\n        for batch in itertoolz.partition_all(bs, files):\n            batch_predictions = onnxinference.predict_batch(batch, bs)\n            all_preds.append(batch_predictions.batch_labels)\n            predictions.append(batch_predictions)\n            progress.update(len(batch))\n            write_batch_preds_to_csv(csv_fname, batch_predictions)\n        all_preds = list(itertoolz.concat(all_preds))\n    typer.echo(f\"CSV report stored in {csv_fname}\")\n    delta = timedelta(seconds=time.perf_counter() - start_time)\n    typer.echo(f\"Time taken to run:  {str(delta)}\")\n    print_table(all_preds)\n\n\ndef print_table(decoded) -> None:\n    \"\"\"Prints table summary of predicted labels\"\"\"\n    table = Table(show_header=True, title=\"Prediction summary\")\n    table.add_column(\n        \"Class\",\n    )\n    table.add_column(\"Count\")\n    table.add_column(\"Percentage\")\n    total = len(decoded)\n    frequencies = itertoolz.frequencies(decoded)\n    for is_last_element, var in core.signal_last(frequencies.items()):\n        key, value = var\n        count = value\n        percentage = round((count / total) * 100, 2)\n        if is_last_element:\n            table.add_row(key, str(count), f\"{percentage}\", end_section=True)\n            table.add_row(\"Total\", str(total), \"\")\n        else:\n            table.add_row(key, str(count), f\"{percentage}\")\n    console.print(table)\n\n\ndef create_csv_fname(csv_directory: Path) -> Path:\n    \"\"\"Creates a csv filename\"\"\"\n    date_now = datetime.now()\n    date_now = date_now.strftime(\"%Y_%m_%d_%H_%M\")\n    fname = Path(date_now + \".csv\")\n    return Path(csv_directory / fname)\n\n\ndef create_csv_header(csv_path: Path) -> None:\n    \"\"\"Creates a header for csv `csv_path`\"\"\"\n    with open(csv_path, mode=\"w\", newline=\"\") as csv_file:\n        field_names = [\"path\", \"directory\", \"predicted_label\", \"confidence\"]\n        writer = csv.DictWriter(csv_file, fieldnames=field_names)\n        writer.writeheader()\n\n\ndef write_batch_preds_to_csv(csv_fpath: Path, predictions: PredictionBatch) -> None:\n    \"\"\"Appends `predictions` batch to `csv_path`\"\"\"\n    with open(csv_fpath, mode=\"a\", newline=\"\") as csv_file:\n        field_names = [\"path\", \"directory\", \"predicted_label\", \"confidence\"]\n        writer = csv.DictWriter(csv_file, fieldnames=field_names)\n        for pred in predictions.batch:\n            row = asdict(pred)\n            row[\"directory\"] = pred.path.parent\n            writer.writerow(row)\n\n\nclass InferenceSession(ABC):\n    \"\"\"Abstract class for inference sessions\"\"\"\n\n    @abstractmethod\n    def __init__(self, model: Path, vocab: List):\n        \"\"\"Inference Sessions should init from a model file and vocab\"\"\"\n        self.model = model\n        self.vocab = vocab\n\n    @abstractmethod\n    def predict_image(self, image: Path):\n        \"\"\"Predict a single image\"\"\"\n        pass\n\n    @abstractmethod\n    def predict_batch(self, model: Path, batch: Iterable[Path], bs: int):\n        \"\"\"Predict a batch\"\"\"\n        pass\n\n\ndef softmax(x):\n    \"\"\"return softmax of `x`\"\"\"\n    x = x.reshape(-1)\n    e_x = np.exp(x - np.max(x))\n    return e_x / e_x.sum(axis=0)\n\n\n# class FastaiInferenceModel(InferenceSession):\n#     def __init__(self, model):\n#         self.model = model\n#         self.learn = load_learner(model)\n\n#     def predict_image(self, image: Path) -> Any:\n#         return self.learn.predict(image)\n\n#     def predict_batch(self, batch: Iterable[Path], bs: int) -> PredictionBatch:\n#         test_dl = self.learn.dls.test_dl(batch, bs=bs)\n#         vocab = dict(enumerate(self.learn.dls.vocab))\n#         with self.learn.no_bar():\n#             fastai_preds: Any = self.learn.get_preds(dl=test_dl, with_decoded=True)\n#             prediction_tensors: Iterable[Any] = fastai_preds[0]\n#             prediction_items = []\n#             for file, pred in zip(batch, prediction_tensors):\n#                 arg_max = int(np.array(pred).argmax())\n#                 predicted_label = vocab[int(arg_max)]\n#                 confidence = float(np.array(pred).max())\n#                 prediction_items.append(\n#                     ImagePredictionItem(file, predicted_label, confidence)\n#                 )\n#         return PredictionBatch(prediction_items)\n\n\nclass OnnxInferenceSession(InferenceSession):\n    \"\"\"onnx inference session\"\"\"\n\n    def __init__(self, model: Path, vocab: Path):\n        \"\"\"Create onnx session\"\"\"\n        self.model = model\n        self.session = rt.InferenceSession(str(model))\n\n        self.vocab = vocab\n        self.vocab_mapping = dict(enumerate(self.vocab))\n\n    def _load_vocab(self, vocab: Path) -> List:\n        with open(vocab, \"r\") as f:\n            return [item.strip(\"\\n\") for item in f.readlines()]\n\n    def predict_image(self, image: Path):\n        \"\"\"Predict a single image\"\"\"\n        img = self._load_image(image)\n        raw_result = self.session.run([\"output\"], {\"image\": img})\n        pred = self._postprocess(raw_result)\n        arg_max = int(np.array(pred).argmax())\n        predicted_label = self.vocab_mapping[int(arg_max)]\n        confidence = float(np.array(pred).max())\n        return ImagePredictionItem(image, predicted_label, confidence)\n\n    def _preprocess(self, input_data: np.ndarray) -> np.ndarray:\n        # converts the input data into the float32 input for onnx\n        img_data = input_data.astype(\"float32\")\n\n        # normalize\n        mean_vec = np.array([0.485, 0.456, 0.406])\n        stddev_vec = np.array([0.229, 0.224, 0.225])\n        norm_img_data = np.zeros(img_data.shape).astype(\"float32\")\n        for i in range(img_data.shape[0]):\n            norm_img_data[i, :, :] = (\n                img_data[i, :, :] / 255 - mean_vec[i]\n            ) / stddev_vec[i]\n\n        # add batch channel\n        norm_img_data = norm_img_data.reshape(1, 3, 512, 512).astype(\"float32\")\n        return norm_img_data\n\n    def _load_image(self, file: Path) -> np.ndarray:\n        \"\"\"loads image and carries out preprocessing for inference\"\"\"\n        image = Image.open(file, mode=\"r\")\n        image = image.resize((512, 512), Image.BILINEAR)\n        image_data = np.array(image).transpose(2, 0, 1)\n        return self._preprocess(image_data)\n\n    def _postprocess(self, result: List):\n        \"\"\"process results from onnx session\"\"\"\n        return softmax(np.array(result)).tolist()\n\n    def predict_batch(self, batch: Iterable[Path], bs: int):\n        \"\"\"predicts a batch of images\"\"\"\n        prediction_items = [self.predict_image(file) for file in batch]\n        return PredictionBatch(prediction_items)\n\n\nif __name__ == \"__main__\":\n    app()\n", "path": "src/flyswot/inference.py"}]}