Demonstration of WhatColorIsX python module, using local image files to make a compound image showing: source image(s), average colour(s), bright hue colour(s)

whatcolorisx_demo.py

import os
import sys
import glob
import argparse
from itertools import zip_longest


import WhatColorIsX
from PIL import Image, ImageDraw, ImageColor, ImageOps, ImageFont


def main(strings=[], strings_file=None, image_dir=None, width=200, output_file='WhatColorIsX_comparison.jpg'):
    """Produces a visualisation of the WhatColorIsX module using image files and/or strings.
    """
    
    # Colours to output
    BRIGHT_HUES = [False, True]
    METHODS = ['average_color'] # ,'common_color']
    
    WIDTH = int(width)
    EXT = ['jpg', 'png']
    COLOR_WIDTH = int(WIDTH/5)
    colour_coord = WIDTH-COLOR_WIDTH
    BORDER = int(WIDTH/200)
    ROW_HEIGHT = int(colour_coord/1.5)
    FONT_SIZE = int(ROW_HEIGHT/6)
    BG_COLOR = (255, 255, 255)
    STRING_COLOR = (0, 0, 0)
    
    try:
        FONT = ImageFont.truetype(size=FONT_SIZE, font='SourceCodePro-Regular.otf')
    except OSError:
        FONT = ImageFont.load_default()
    
    # Get files with one of EXT as extension
    files = []
    if image_dir:
        for e in EXT:
            files.extend(glob.glob(os.path.join(image_dir, '*.{0}'.format(e))))
    
    # Get list of strings from file
    if strings_file:
        with open(strings_file, 'r') as fh:
            strings.extend(fh.read().splitlines())
    
    # Intercalate both lists together
    things = [y for x in zip_longest(files, strings) for y in x if y is not None]
    num_things = len(files)+len(strings)

    # If nothing found, exit
    if num_things:
        canvas = Image.new('RGB', (WIDTH, ROW_HEIGHT*num_things), BG_COLOR)
        canvas_draw = ImageDraw.Draw(canvas)
        # For each image
        for i, thing in enumerate(things):
            print('Processing {0} of {1}: {2}'.format(i+1, num_things, thing))
            img_crop_size = (colour_coord-BORDER*2, ROW_HEIGHT-BORDER*2)
            colors = []
            # If image file
            try:
                img = Image.open(thing)
                to_wcix = img
                # Fit and paste image into collage
                img_crop = ImageOps.fit(img, img_crop_size, method=Image.BICUBIC)
            # Else treat as string, search and print in image
            except FileNotFoundError:
                to_wcix = thing
                img_crop = Image.new('RGB', img_crop_size, BG_COLOR)
                img_crop_draw = ImageDraw.Draw(img_crop)
                string_dimensions = FONT.getsize(thing)
                string_topleft = ([int(img_crop.size[x]*0.5-string_dimensions[x]*0.5) for x in range(2)])
                img_crop_draw.text(string_topleft, thing, font=FONT, fill=STRING_COLOR)
                
            # Find colours of img
            for m in METHODS:
                for bh in BRIGHT_HUES:
                    wcix = WhatColorIsX.new(to_wcix)
                    colors.append(wcix.color(method=m, bright_hue=bh))
            
            # Paste img_crop to canvas
            canvas.paste(img_crop, (BORDER, i*ROW_HEIGHT+BORDER))
            
            # Add colour swatches
            swatch_multiplier = 1/len(colors)
            for j, color in enumerate(colors):
                coords = (colour_coord+BORDER, (i+j*swatch_multiplier)*ROW_HEIGHT+BORDER, WIDTH-BORDER, (i+(j+1)*swatch_multiplier)*ROW_HEIGHT-BORDER-1) # Not sure why -1 here, think rectangle draws incorrectly
                canvas_draw.rectangle(coords, fill=ImageColor.getrgb(color))                
                
        print('Saving...')
        canvas.save(output_file)
    else:
        sys.exit('No valid files/strings provided')


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Demo functionality of WhatColorIsX module')
    parser.add_argument('-S', '--string', nargs='+', default=[],
                        help='Strings to search for')
    parser.add_argument('-s', '--strings_file',
                        help='Source file for newline-separated strings to search')
    parser.add_argument('-i', '--image_dir',
                        help='Source directory for local image files')
    parser.add_argument('-w', '--width', default=200,
                        help='Width of the output image')
    parser.add_argument('-o', '--output_file', default='WhatColorIsX_demo.png',
                        help='Output path for the image file')
    args = parser.parse_args()
    main(strings=args.string, strings_file=args.strings_file, image_dir=args.image_dir, width=args.width, output_file=args.output_file)