pai2/utils/watermark_utils.py

# -*- coding: utf-8 -*-

try:
    import Image
    import ImageEnhance
except ImportError:
    from PIL import Image, ImageEnhance


def reduce_opacity(im, opacity):
    """Returns an image with reduced opacity."""
    assert 0 <= opacity <= 1
    im = im.convert('RGBA') if im.mode != 'RGBA' else im.copy()
    alpha = im.split()[3]
    alpha = ImageEnhance.Brightness(alpha).enhance(opacity)
    im.putalpha(alpha)
    return im


def watermark(im, mark, position, opacity=1):
    """Adds a watermark to an image."""
    if opacity < 1:
        mark = reduce_opacity(mark, opacity)
    if im.mode != 'RGBA':
        im = im.convert('RGBA')
    # create a transparent layer the size of the image and draw the
    # watermark in that layer.
    layer = Image.new('RGBA', im.size, (0, 0, 0, 0))
    if position == 'tile':
        for y in range(0, im.size[1], mark.size[1]):
            for x in range(0, im.size[0], mark.size[0]):
                layer.paste(mark, (x, y))
    elif position == 'scale':
        # scale, but preserve the aspect ratio
        ratio = min(
            float(im.size[0]) / mark.size[0], float(im.size[1]) / mark.size[1])
        w = int(mark.size[0] * ratio)
        h = int(mark.size[1] * ratio)
        mark = mark.resize((w, h))
        layer.paste(mark, ((im.size[0] - w) / 2, (im.size[1] - h) / 2))
    else:
        layer.paste(mark, position)
    # composite the watermark with the layer
    return Image.composite(layer, im, layer)


def watermark_wrap(im_path, mark_path, save_path=''):
    im, mark = Image.open(im_path), Image.open(mark_path)
    new_im = watermark(im, mark, (50, 50), 0.5)
    new_im.save(save_path or im_path)


def watermark_test():
    im, mark = Image.open('original_CGzC_10a50000c8811190.jpg'), Image.open('paiai_96_96.png')
    watermark(im, mark, 'tile', 0.5).show()
    watermark(im, mark, 'scale', 1.0).show()
    watermark(im, mark, (50, 50), 0.5).show()


if __name__ == '__main__':
    # watermark_test()
    watermark_wrap('original_CGzC_10a50000c8811190.jpg', 'paiai_96_96.png')