Extract NPZ file which contains images as torch compatible arrays and save as actual image files

npz_to_image.py

"""A python script to convert Numpy NPZ files containing torch-compatible shaped arrays of images

You can run this file with command:
python npz_to_to_image.py [PATH_TO_NPZ_FILE] [PATH_TO_DESTINATION_FOLDER] [IMAGE_EXTENSION]

"""
import os
import numpy as np
import argparse
from tqdm import tqdm
from PIL import Image


def main(npz_fp : str, dest_path : str, image_extension : str):
    """Saves all images saved in a NPZ as arrays to files

    Args:
        npz_fp (str): Absolute or relative path to NPZ file
        dest_path (str): Absolute or relative to destination directory
    """    
    # Load NPZ file containing images
    with np.load(os.path.abspath(npz_fp), allow_pickle=True) as npz:

        # Create destination directory
        dest_full_path = os.path.abspath(dest_path)
        
        if not os.path.exists(dest_full_path):
            os.makedirs(dest_full_path)

        for npy in tqdm(npz, desc="Saving images", unit=" Image"):

            # Get array of image in torch compatible shape and reshape it
            pil_image = Image.fromarray(torch_compatible_array_to_pil(npz[npy]))
            # Remove .npz from file name
            image_name = (explode_string(npy, '.'))[0]
            pil_image.save(os.path.join(dest_full_path, f"{image_name}.{image_extension}"))


def explode_string(string : str, divider : str) -> list:
    """Return array of words divided by divider

    Args:
        string (str): The string to explode into array
        divider (str):The character to divide string with

    Returns:
        list: A list of all words from string
    """    
    return string.split(divider)


def torch_compatible_array_to_pil(array):
    """Return image array in the accepted shape for PIL

    Args:
        array (nd.array): Array of image with shape (c, w, h)

    Returns:
        nd.array: Array of image with shape (w, h, c)
    """    
    return (np.transpose(array, (1, 2, 0))).astype("uint8")


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("npz_fp", type=str,
                        help="Path to NPZ file")
    parser.add_argument("dest_path", type=str,
                        help="Destination path to image files")
    parser.add_argument("img_extension", type=str,
                        help="Images extension")
    args = parser.parse_args()
    main(args.npz_fp, args.dest_path, args.img_extension)