image.py 文件源码

def array_to_image(
        color_images_array=None,
        gray_images_array=None,
        mode='RGB',
        color_normalize=False,
        linedrawing=None,
):
    # type: (any,any,any,any,any) -> typing.List[Image.Image]
    """
    :param color_images_array: shape is [number of image, channel(3), width, height]
    :param gray_images_array: used when mode=='ab' or 'gray'
    :param mode: mode of input images array (RGB, Lab, ab, gray)
    :param color_normalize: normalize rgb color to [min(rgb_images_array) max(rgb_images_array)]
    """
    if color_images_array is not None:
        color_images_array = chainer.cuda.to_cpu(color_images_array)

    if gray_images_array is not None:
        gray_images_array = chainer.cuda.to_cpu(gray_images_array)

    if mode == 'gray':
        color_images_array = numpy.concatenate([gray_images_array] * 3, axis=1)
        mode = 'RGB'

    if mode == 'ab':
        # concat gray image(luminance) and ab(chromaticity)
        color_images_array = chainer.cuda.to_cpu(color_images_array)
        color_images_array = numpy.concatenate((gray_images_array, color_images_array), axis=1)
        mode = 'Lab'

    color_images_array = color_images_array.transpose(0, 2, 3, 1)

    if mode == 'Lab':
        color_images_array = color_images_array.astype(dtype=numpy.float64)
        image_array_list = [lab2rgb(image_array) * 255 for image_array in color_images_array]
        color_images_array = numpy.concatenate(
            [numpy.expand_dims(image_array, axis=0) for image_array in image_array_list]
        )
        mode = 'RGB'

    if mode == 'RGB':
        rgb_images_array = color_images_array
    else:
        raise ValueError('{} mode is not supported'.format(mode))

    # to uint8
    if color_normalize:
        minmax = (rgb_images_array.min(), rgb_images_array.max())
    else:
        if linedrawing is not None:
            minmax = (0, 1)
        else:
            minmax = (0, 255)

    def clip_image(x):
        x = (x - minmax[0]) / (minmax[1] - minmax[0]) * 255  # normalize to 0~255
        return numpy.float32(0 if x < 0 else (255 if x > 255 else x))

    rgb_images_array = numpy.vectorize(clip_image)(rgb_images_array)
    rgb_images_array = rgb_images_array.astype(numpy.uint8)
    return [Image.fromarray(image_array) for image_array in rgb_images_array]