lettersim-ot-rsa/scold/draw.py

"""Functions for drawing text"""

from PIL import Image, ImageDraw, ImageFont
import numpy as np
from string import ascii_letters
from scipy.signal import convolve2d

def text_array(text, font='arial.ttf', size=50, colour=(255,255,255), bg=(0,0,0), border=(0,0,0,0), crop_x='text', crop_y='text', crop_chars=ascii_letters, fontcrop_n=None, align_x='centre', greyscale=True, method=np.round, rotate=0, fliplr=False, flipud=False, outline=False, return_img=False):
    """Return an `np.array` (default) or PIL image of the specified text, cropped to its boundaries.
    
    Paramaters
    ----------
    text : str
        The text to draw.
    font : str
        The .ttf font file to use.
    size : int
        The font size in number of pixels (passed to `PIL.ImageFont.truetype()`)
    colour : tuple
        A tuple of length 3, specifying the font colour (0 to 255) in in the form (R,G,B)
    bg : tuple
        A tuple specifying the background colour (0 to 255) in in the form (R,G,B)
    border : tuple
        A tuple of pixel sizes for a border in the form (left_x, right_x, top_y, bottom_y), surrounding the crop
    crop_x : str
        Should be one of the following, specifiying whether the output's width be cropped to the maximum boundaries of the current text or the font:
        'text': the width limits of the rendered text
        'font': the width limits of the font's ascii characters given the text's length, with centred alignment
    crop_y : str
        Should be one of the following, specifiying whether the output's height be cropped to the maximum boundaries of the current text or the font:
        'text': the height limits of the rendered text
        'font': the height limits of the font's ascii characters given the text's length
    crop_chars : list
        A list of characters that should be used to find the font extremities when cropping. Default is `string.ascii_letters`.
    fontcrop_n : int
        If crop_x or crop_y are 'font', how many characters should the crop assume are available to the font? If `None` (default) will use `len(text)`.
    align_x : str
        How to horizontally align the text (if any white space):
        'left': align to the left
        'centre': align to the centre
        'right': align to the right
    greyscale : bool
        Should the resulting image be converted to greyscale?
    method : function, optional
        Function to apply to the array at the end. Default is `np.round` to binarise.
    rotate : float
        Degrees by which the text should be rotated.
    fliplr : bool
        Should the array be mirrored horizontally?
    flipud : bool
        Should the array be mirrored vertically?
    outline : bool
        Should the array be outlined using `outline_shape()`?
    return_img : bool
        Should an image be returned, rather than an array?
    """
    
    # get font info
    pil_font = ImageFont.truetype(font, size=size, encoding="unic")
    pil_fontsize = pil_font.getbbox(text)[-2:]
    
    # draw image
    im = Image.new('RGB', pil_fontsize, bg)
    im_draw = ImageDraw.Draw(im)
    im_draw.text((0,0), text, font=pil_font, fill=colour)
    
    # find which pixels could be filled by text
    if crop_x=='font' or crop_y=='font':
        if fontcrop_n is None:
            crop_nchar = len(text)
        else:
            crop_nchar = fontcrop_n
        letter_fontsizes = [pil_font.getbbox(letter*crop_nchar)[-2:] for letter in set(crop_chars)]
        canvas_max_size = (max([tup[0] for tup in letter_fontsizes]), max([tup[1] for tup in letter_fontsizes]))
        im_lims = Image.new('RGB', canvas_max_size, bg)
        im_lims_draw = ImageDraw.Draw(im_lims)
        for letter in crop_chars:
            im_lims_draw.text((0,0), letter*crop_nchar, font=pil_font, fill=colour)
        im_lims_arr = np.array(im_lims)
        text_lims_px = np.where(np.sum(im_lims_arr==bg, 2)!=im_lims_arr.shape[2])
        
    # find which pixels are filled with text
    im_arr = np.array(im)
    text_px = np.where(np.sum(im_arr==bg, 2)!=im_arr.shape[2])  # find non-background pixels   
    
    # get x crop from either font or text limits
    if crop_x=='font':
        min_x = np.min(text_lims_px[1])
        max_x = np.max(text_lims_px[1])+1
    elif crop_x=='text':
        min_x = np.min(text_px[1])
        max_x = np.max(text_px[1])+1
    
    # get y crop from either font or text limits
    if crop_y=='font':
        min_y = np.min(text_lims_px[0])
        max_y = np.max(text_lims_px[0])+1
    elif crop_y=='text':
        min_y = np.min(text_px[0])
        max_y = np.max(text_px[0])+1
    
    min_x = int(min_x)
    max_x = int(max_x)
    min_y = int(min_y)
    max_y = int(max_y)
    
    # get the horizontal and vertical text size info (for alignment & drawing location)
    min_x_text = np.min(text_px[1])
    max_x_text = np.max(text_px[1])+1
    # min_y_text = np.min(text_px[0])
    # max_y_text = np.max(text_px[0])+1
    text_width = max_x_text - min_x_text
    
    # apply crop (create as new image to avoid default black background)
    im_out = Image.new('RGB', (border[0] + border[1] + max_x - min_x, border[2] + border[3] + max_y - min_y), bg)
    im_out_draw = ImageDraw.Draw(im_out)
    
    # get the text position adjustment (dictated by alignment)
    if align_x=='centre':
        align_adjust = round((max_x - min_x) * 0.5 - text_width * 0.5)
    elif align_x=='left':
        align_adjust = 0
    elif align_x=='right':
        align_adjust = max_x - text_width
    
    # draw the text in the necessary position
    im_out_draw.text((-min_x_text + border[0] + align_adjust, -min_y + border[2]), text, font=pil_font, fill=colour)
    
    # convert to greyscale if necessary
    if greyscale:
        im_out=im_out.convert('L')
    
    # rotate if necessary
    if rotate!=0:
        im_out=im_out.rotate(rotate, expand=True)
    
    # return the PIL image if requested
    if return_img and method is None and not outline:
        return(im_out)
    
    # convert to array
    arr = np.asarray(im_out)
    arr = arr / np.max(arr)

    if fliplr:
        arr = np.fliplr(arr)

    if flipud:
        arr = np.flipud(arr)

    # apply function if required
    if method is not None:
        arr = method(arr)

    # outline the shape if requested
    if outline:
        arr = outline_shape(arr)
    
    # return the PIL image if requested
    if return_img:
        im_out = Image.fromarray(np.uint8(arr*255))
        return(im_out)
        
    return(arr)

def outline_shape(arr, binarise=True, thresh=0):
    """Keep only the outline of non-zero elements of a shape. The outline is located using convolution with kernels for finding the vertical and horizontal edges. The output will be binary.
    
    Paramaters
    ----------
    arr : np.ndarray
        Array to outline - should be scaled between 0 and 1.
    binarise : bool
        Should the array be binarised before finding the outline.
    thresh : float
        Threshold (from 0 to 1) for binarisation.
    """
    arr_bin = np.zeros(arr.shape)
    arr_bin[arr>thresh] = 1

    edge_kernels = [
        # vertical edges
        [[-1, 1, 0]],
        [[0, 1, -1]],
        # horizontal edges
        [[-1], [1], [0]],
        [[0], [1], [-1]]
    ]

    edges = np.zeros(arr.shape)

    for k in edge_kernels:
        c = convolve2d(arr_bin, k, mode='same')
        c[c<0] = 0
        edges += c

    edges[edges>1] = 1  # account for some pixels being edges with multiple sides facing zeroes

    if not binarise:
        # if binarise is False, return the original values in the locations of the edges (useful for methods using weighting)
        edges[edges==1] = arr[edges==1]
    
    return(edges)