jacebrowning / memegen

memegen/domain/image.py

import os
import hashlib
from contextlib import suppress
from typing import Tuple

from PIL import Image as ImageFile, ImageFont, ImageDraw, ImageFilter

Unable to import 'PIL'

import log


MAXIMUM_PIXELS = 1920 * 1080


class Image:
    """JPEG generated by applying text to a template."""

    def __init__(self, template, text, root=None,
                 style=None, font=None, size=None,
                 watermark="", watermark_font=None):
        self.root = root
        self.template = template
        self.style = style
        self.text = text
        self.font = font
        self.width = size.get('width') if size else None
        self.height = size.get('height') if size else None
        self.watermark = watermark
        self.watermark_font = watermark_font

    @property
    def path(self):
        if not self.root:
            return None

        base = os.path.join(self.root, self.template.key, self.text.path)
        custom = [self.style, self.font, self.watermark,
                  self.width, self.height]

        if any(custom):
            slug = self.hash(custom)
            return "{}#{}.img".format(base, slug)
        else:
            return base + ".img"

    @staticmethod
    def hash(values):
        sha = hashlib.md5()
        for index, value in enumerate(values):
            sha.update("{}:{}".format(index, value or "").encode('utf-8'))
        return sha.hexdigest()

    def save(self):
        data = _generate(
            top=self.text.top, bottom=self.text.bottom,
            font_path=self.font.path,
            background=self.template.get_path(self.style),
            width=self.width, height=self.height,
            watermark=self.watermark,
            watermark_font_path=self.watermark_font.path,
        )

        directory = os.path.dirname(self.path)
        with suppress(FileExistsError):
            os.makedirs(directory)

        log.info("Saving image: %s", self.path)
        path = data.save(self.path, format=data.format)

        return path


def _generate(top, bottom, font_path, background, width, height,
              watermark, watermark_font_path):
    """Add text to an image and save it."""
    log.info("Loading background: %s", background)
    background_image = ImageFile.open(background)
    if background_image.mode not in ('RGB', 'RGBA'):
        if background_image.format == 'JPEG':
            background_image = background_image.convert('RGB')
            background_image.format = 'JPEG'
        else:
            background_image = background_image.convert('RGBA')
            background_image.format = 'PNG'

    # Resize to a maximum height and width
    ratio = background_image.size[0] / background_image.size[1]
    pad_image = bool(width and height)
    pad_watermark = True
    if pad_image:
        if width < height * ratio:
            dimensions = width, int(width / ratio)
            pad_watermark = False
        else:
            dimensions = int(height * ratio), height
    elif width:
        dimensions = width, int(width / ratio)
    elif height:
        dimensions = int(height * ratio), height
    else:
        dimensions = 600, int(600 / ratio)
    dimensions = _fit_image(*dimensions)
    image = background_image.resize(dimensions, ImageFile.LANCZOS)
    image.format = 'PNG'

    # Draw image
    draw = ImageDraw.Draw(image)

    max_font_size = int(image.size[1] / 9)
    min_font_size_single_line = int(image.size[1] / 12)
    max_text_len = image.size[0] - 20
    top_font_size, top = _optimize_font_size(
        font_path, top, max_font_size,
        min_font_size_single_line, max_text_len,
    )
    bottom_font_size, bottom = _optimize_font_size(
        font_path, bottom, max_font_size,
        min_font_size_single_line, max_text_len,
    )

    top_font = ImageFont.truetype(font_path, top_font_size)
    bottom_font = ImageFont.truetype(font_path, bottom_font_size)

    top_text_size = draw.multiline_textsize(top, top_font)
    bottom_text_size = draw.multiline_textsize(bottom, bottom_font)

    # Find top centered position for top text
    top_text_pos_x = (image.size[0] / 2) - (top_text_size[0] / 2)
    top_text_pos_y = 0
    top_text_pos = (top_text_pos_x, top_text_pos_y)

    # Find bottom centered position for bottom text
    bottom_text_size_x = (image.size[0] / 2) - (bottom_text_size[0] / 2)
    bottom_text_size_y = image.size[1] - bottom_text_size[1] * (7 / 6)
    if watermark and pad_watermark:
        bottom_text_size_y = bottom_text_size_y - 5
    bottom_text_pos = (bottom_text_size_x, bottom_text_size_y)

    outline = int(round(max(1, min(dimensions) / 300)))
    t_outline = min(outline, top_font_size // 20)
    b_outline = min(outline, bottom_font_size // 20)
    _draw_outlined_text(image, top_text_pos, top, top_font, t_outline)
    _draw_outlined_text(image, bottom_text_pos, bottom, bottom_font, b_outline)

    # Pad image if a specific dimension is requested
    if pad_image:
        image = _add_blurred_background(image, background_image, width, height)

    # Add watermark
    if watermark:
        watermark_font = ImageFont.truetype(watermark_font_path, 11)
        watermark_pos = (3, image.size[1] - 15)
        _draw_outlined_text(image, watermark_pos, watermark, watermark_font,
                            alpha=True)

    return image


def _optimize_font_size(font, text, max_font_size, min_font_size,
                        max_text_len):
    """Calculate the optimal font size to fit text in a given size."""

    # Check size when using smallest single line font size
    fontobj = ImageFont.truetype(font, min_font_size)
    text_size = fontobj.getsize(text)

    # Calculate font size for text, split if necessary
    if text_size[0] > max_text_len:
        phrases = _split(text)
    else:
        phrases = (text,)
    font_size = max_font_size
    for phrase in phrases:
        font_size = min(_maximize_font_size(font, phrase, max_text_len),
                        font_size)

    # Rebuild text with new lines
    text = '\n'.join(phrases)

    return font_size, text


def _draw_outlined_text(image, text_pos, text, font, width=1, alpha=False):
    """Draw white text with black outline on an image."""
    if alpha:
        black_alpha = 170
        white_alpha = 128
    else:
        black_alpha = 255
        white_alpha = 240

    overlay = ImageFile.new('RGBA', image.size)
    draw = ImageDraw.ImageDraw(overlay, 'RGBA')

    # Draw black text outlines
    for x in range(-width, 1 + width):
        for y in range(-width, 1 + width):
            pos = (text_pos[0] + x, text_pos[1] + y)
            draw.multiline_text(pos, text, (0, 0, 0, black_alpha),
                                font=font, align='center')

    # Draw inner white text
    draw.multiline_text(text_pos, text, (255, 255, 255, white_alpha),
                        font=font, align='center')

    image.paste(overlay, None, overlay)


def _add_blurred_background(foreground, background, width, height):
    """Add a blurred background to match the requested dimensions."""
    base_width, base_height = foreground.size

    border_width = min(width, base_width + 2)
    border_height = min(height, base_height + 2)
    border_dimensions = border_width, border_height
    border = ImageFile.new('RGB', border_dimensions)
    border.paste(foreground, ((border_width - base_width) // 2,
                              (border_height - base_height) // 2))

    padded_dimensions = _fit_image(width, height)
    padded = background.resize(padded_dimensions, ImageFile.LANCZOS)

    darkened = padded.point(lambda p: p * 0.4)

    blurred = darkened.filter(ImageFilter.GaussianBlur(5))
    blurred.format = 'PNG'

    blurred_width, blurred_height = blurred.size
    offset = ((blurred_width - border_width) // 2,
              (blurred_height - border_height) // 2)
    blurred.paste(border, offset)

    return blurred


def _maximize_font_size(font, text, max_size):
    """Find the biggest font size that will fit."""
    font_size = max_size

    fontobj = ImageFont.truetype(font, font_size)
    text_size = fontobj.getsize(text)
    while text_size[0] > max_size and font_size > 1:
        font_size = font_size - 1
        fontobj = ImageFont.truetype(font, font_size)
        text_size = fontobj.getsize(text)

    return font_size


def _fit_image(width: int, height: int) -> Tuple[int, int]:
    while width * height > MAXIMUM_PIXELS:
        width *= 0.75
        height *= 0.75
    return int(width), int(height)


def _split(text):
    """Split a line of text into two similarly sized pieces.

    >>> _split("Hello, world!")
    ('Hello,', 'world!')

    >>> _split("This is a phrase that can be split.")
    ('This is a phrase', 'that can be split.')

    >>> _split("This_is_a_phrase_that_can_not_be_split.")
    ('This_is_a_phrase_that_can_not_be_split.',)

    """
    result = (text,)

    if len(text) >= 3 and ' ' in text[1:-1]:  # can split this string
        space_indices = [i for i in range(len(text)) if text[i] == ' ']
        space_proximities = [abs(i - len(text) // 2) for i in space_indices]
        for i, j in zip(space_proximities, space_indices):
            if i == min(space_proximities):
                result = (text[:j], text[j + 1:])
                break

    return result