scripts/saver.py

from __future__ import annotations

import numpy as np
import torch
from PIL import Image


class JPGSaver:
    def __init__(self, output_dir: str) -> None:
        super().__init__()
        self.output_dir = output_dir

    def save(self, image_data: torch.Tensor, file_name: str) -> None:
        image_data = np.clip(image_data.cpu().numpy(), 0, 1)
        image_data = (image_data * 255).astype(np.uint8)
        im = Image.fromarray(image_data[0, 0])
        im.save(self.output_dir + "/" + file_name + ".jpg")