Skip to content
Snippets Groups Projects
Commit ef7494b7 authored by Daniel CARRON's avatar Daniel CARRON :b:
Browse files

Added util to save images

parent 4ccba39c
No related branches found
No related tags found
1 merge request!7Reviewed DataModule design+docs+types
# 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)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment