Python script to store nd array into images

Question

I am trying to write a generic Python script that do the following types of tasks:

1. Load .npy file (the .npy file is of shape (m_samples, channels, row, column), which corresponds to m_samples images, and here the channels are always 1
2. Check whether a given path exists
3. Iterate against the .npy file, and save each image into the given path

I am not sure whether I am using the best practices to save an nd array into an image.

import numpy as np
from scipy.misc import toimage, imsave
import os
image_path = "raw"

def ensure_directory_exist(image_path):
    if not os.path.exists(image_path):
        print("Allocating '{:}'",format(image_path))
        os.mkdir(image_path)


if __name__ == '__main__':
    img_array = np.load('imgs_mask_test.npy')
    ensure_directory_exist(image_path)
    for i in range(img_array.shape[0]):
       name = "img"+str(i)+".png"
       imsave(os.path.join(image_path,name),img_array[i,0])

Answer 1

1. Instead of:

   ensure_directory_exist(image_path)
   

   use the built-in os.makedirs:

   os.makedirs(image_path, exist_ok=True)
   

2. Iterate over the images themselves (rather than their indexes) and so avoid the lookup img_array[i, 0] on each iteration:

   for i, img in enumerate(img_array[:, 0, :, :]):
       name = "img{}.png".format(i)
       imsave(os.path.join(image_path, name), img)