#!/usr/bin/env python2

'''Read and write image files as NumPy arrays'''


from numpy import asarray, float32
from PIL import Image

from . import np
from . import utils


_DEFAULT_DTYPE = float32

_PIL_RGB = {
    'R': 0,
    'G': 1,
    'B': 2,
}


def get_channels(img):
    '''Return a list of channels of an image array.'''
    if utils.depth(img) == 1:
        yield img
    else:
        for i in xrange(utils.depth(img)):
            yield img[:, :, i]


def read(filename, dtype=_DEFAULT_DTYPE):
    '''Read an image file as an array.'''
    img = Image.open(filename)
    arr = asarray(img, dtype=dtype)
    arr = utils.swap_rgb(arr, _PIL_RGB)
    return arr


def _pil_save(img, filename):
    pil_img = Image.fromarray(img)
    pil_img.save(filename)
    return


def save(img, filename, random=False):
    '''Save an image array and return its path.'''
    if random:
        newfile = utils.rand_filename(filename)
    else:
        newfile = filename
    np.normalize(img)
    uint8img = np.to_uint8(img)
    _pil_save(uint8img, newfile)
    return newfile


if __name__ == '__main__':
    pass

#!/usr/bin/env python2

'''Read and write image files as NumPy arrays'''


from numpy import asarray, float32
from PIL import Image

from . import np
from . import utils


_DEFAULT_DTYPE = float32

_PIL_RGB = {
    'R': 0,
    'G': 1,
    'B': 2,
}


def get_channels(img):
    '''Return a list of channels of an image array.'''
    if utils.depth(img) == 1:
        yield img
    else:
        for i in xrange(utils.depth(img)):
            yield img[:, :, i]


def read(filename, dtype=_DEFAULT_DTYPE):
    '''Read an image file as an array.'''
    img = Image.open(filename)
    arr = asarray(img, dtype=dtype)
    arr = utils.swap_rgb(arr, _PIL_RGB)
    return arr


def _pil_save(img, filename):
    pil_img = Image.fromarray(img)
    pil_img.save(filename)
    return


def save(img, filename, random=False, ext=None):
    '''Save an image array and return its path.'''
    if random:
        newfile = utils.rand_filename(filename, ext=ext)
    else:
        newfile = filename
    np.normalize(img)
    uint8img = np.to_uint8(img)
    _pil_save(uint8img, newfile)
    return newfile


if __name__ == '__main__':
    pass

Add an extension parameter to the save function.