Python script to store nd array into images
Instead of:

ensure_directory_exist(image_path)


use the built-in os.makedirs:

os.makedirs(image_path, exist_ok=True)

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)