From 0ac06e95d652dbaa89edba9d5b4179cf939529e2 Mon Sep 17 00:00:00 2001
From: dcarron <daniel.carron@idiap.ch>
Date: Wed, 5 Jul 2023 13:34:02 +0200
Subject: [PATCH] Added util to save images

---
 src/ptbench/utils/image.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)
 create mode 100644 src/ptbench/utils/image.py

diff --git a/src/ptbench/utils/image.py b/src/ptbench/utils/image.py
new file mode 100644
index 00000000..363a8309
--- /dev/null
+++ b/src/ptbench/utils/image.py
@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright © 2023 Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import os
+
+from typing import Union
+
+import torch
+
+from PIL.Image import Image
+from torchvision import transforms
+
+
+def save_image(img: Union[torch.Tensor, Image], filepath: str) -> None:
+    """Saves a PIL image or a tensor as an image at the specified destination.
+
+    Parameters
+    ----------
+
+    img:
+        A torch.Tensor or PIL.Image to save
+
+    filepath:
+        The file in which to save the image. The format is inferred from the file extension, or defaults to png if not specified.
+    """
+
+    if isinstance(img, torch.Tensor):
+        img = transforms.ToPILImage()(img)
+
+    root, ext = os.path.splitext(filepath)
+
+    if len(ext) == 0:
+        filepath = filepath + ".png"
+
+    img.save(filepath)
-- 
GitLab