We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
check csv output doesn't have blank lines 


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

{"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@@ -112,15 +112,14 @@\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 Progress() as progress:\n-        prediction_progress = progress.add_task(\"Predicting images\", total=len(files))\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(prediction_progress, advance=bs)\n+            progress.update(len(batch))\n             write_batch_preds_to_csv(csv_fname, batch_predictions)\n         all_preds = list(itertoolz.concat(all_preds))\n     print_table(all_preds)\n@@ -155,14 +154,14 @@\n \n \n def create_csv_header(csv_path: Path) -> None:\n-    with open(csv_path, mode=\"w\") as csv_file:\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 \n def write_batch_preds_to_csv(csv_fpath: Path, predictions: PredictionBatch) -> None:\n-    with open(csv_fpath, mode=\"a\") as csv_file:\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", "issue": "check csv output doesn't have blank lines \n\n", "before_files": [{"content": "\"\"\"Inference functionality\"\"\"\nimport csv\nimport mimetypes\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom dataclasses import asdict\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom pathlib import Path\nfrom typing import IO\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\nimport PIL\nimport typer\nfrom PIL import Image\nfrom rich.progress import Progress\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\ntry:\n    from fastai.vision.all import Learner\n    from fastai.vision.all import load_learner\nexcept ImportError:\n    pass\nfrom importlib import resources\n\napp = typer.Typer()\n\n\n# flake8: noqa\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        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        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    pass  # pragma: no cover\n\n\n@app.command(name=\"directory\")\ndef predict_directory(\n    directory: Path = typer.Argument(..., readable=True, resolve_path=True),\n    csv_save_dir: Path = typer.Argument(\n        ...,\n        writable=True,\n        resolve_path=True,\n    ),\n    pattern: str = typer.Option(\"fse\"),\n    bs: int = typer.Option(32),\n    preferred_format: str = typer.Option(\n        \".tif\",\n        help=\"Preferred image format for predictions. If not available, flyswot will use images matching `pattern`\",\n    ),\n):\n    \"\"\"\n    Predicts against all images containing PATTERN in the filename found under DIRECTORY.\n    By default searches for filenames containing FSE\n    Creates a CSV report saved to CSV_SAVE_DIR containing the predictions\n    \"\"\"\n    typer.echo(csv_save_dir)\n    model_dir = models.ensure_model_dir()\n    typer.echo(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 = onnx_inference_session(model, vocab)\n    files = core.get_image_files_from_pattern(directory, pattern)\n    filtered_files = core.filter_to_preferred_ext(files, preferred_format)\n    files = list(filtered_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    print_table(all_preds)\n\n\ndef print_table(decoded) -> None:\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    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    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    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    @abstractmethod\n    def __init__(self, model: Path, vocab: List):\n        self.model = model\n        self.vocab = vocab\n\n    @abstractmethod\n    def predict_image(self, image: Path):\n        pass\n\n    @abstractmethod\n    def predict_batch(self, model: Path, batch: Iterable[Path], bs: int):\n        pass\n\n\ndef softmax(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 onnx_inference_session(InferenceSession):\n    \"\"\"Class for running inference making use of the onnxrunntime\"\"\"\n\n    def __init__(self, model, vocab):\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\nfrom abc import ABC\nfrom abc import abstractmethod\nfrom dataclasses import asdict\nfrom dataclasses import dataclass\nfrom datetime import datetime\nfrom pathlib import Path\nfrom typing import IO\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\nimport PIL\nimport typer\nfrom PIL import Image\nfrom rich.progress import Progress\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\ntry:\n    from fastai.vision.all import Learner\n    from fastai.vision.all import load_learner\nexcept ImportError:\n    pass\nfrom importlib import resources\n\napp = typer.Typer()\n\n\n# flake8: noqa\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        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        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    pass  # pragma: no cover\n\n\n@app.command(name=\"directory\")\ndef predict_directory(\n    directory: Path = typer.Argument(..., readable=True, resolve_path=True),\n    csv_save_dir: Path = typer.Argument(\n        ...,\n        writable=True,\n        resolve_path=True,\n    ),\n    pattern: str = typer.Option(\"fse\"),\n    bs: int = typer.Option(32),\n    preferred_format: str = typer.Option(\n        \".tif\",\n        help=\"Preferred image format for predictions. If not available, flyswot will use images matching `pattern`\",\n    ),\n):\n    \"\"\"\n    Predicts against all images containing PATTERN in the filename found under DIRECTORY.\n    By default searches for filenames containing FSE\n    Creates a CSV report saved to CSV_SAVE_DIR containing the predictions\n    \"\"\"\n    typer.echo(csv_save_dir)\n    model_dir = models.ensure_model_dir()\n    typer.echo(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 = onnx_inference_session(model, vocab)\n    files = core.get_image_files_from_pattern(directory, pattern)\n    filtered_files = core.filter_to_preferred_ext(files, preferred_format)\n    files = list(filtered_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    print_table(all_preds)\n\n\ndef print_table(decoded) -> None:\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    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    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    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    @abstractmethod\n    def __init__(self, model: Path, vocab: List):\n        self.model = model\n        self.vocab = vocab\n\n    @abstractmethod\n    def predict_image(self, image: Path):\n        pass\n\n    @abstractmethod\n    def predict_batch(self, model: Path, batch: Iterable[Path], bs: int):\n        pass\n\n\ndef softmax(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 onnx_inference_session(InferenceSession):\n    \"\"\"Class for running inference making use of the onnxrunntime\"\"\"\n\n    def __init__(self, model, vocab):\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"}]}