Automatically extract diecut paths for an image.

make_diecut.py

#!/usr/bin/env python2

from __future__ import print_function
import argparse
import re
import sys

import attr
import cv2
import numpy
from PIL import Image
import scipy.interpolate

try:
    import ipdb
except ImportError:
    pass # For debugging only.


@attr.s
class RunState(object):
    args = attr.ib()
    base = attr.ib()
    key = attr.ib()
    outline = attr.ib()
    mask = attr.ib()
    diecut = attr.ib()


def get_key_pixel(state, x, y):
    key_x = x % state.key.width
    key_y = y % state.key.height
    return state.key.getpixel((key_x, key_y))


def outline_blend(base_pixel, key_pixel):
    rgb_diff = sum(abs(b_c - k_c) for (b_c, k_c) in zip(base_pixel, key_pixel))
    if rgb_diff < 30:
        return (0,0,0,0)
    else:
        return base_pixel


def generate_outline(state, x, y):
    base_pixel = state.base.getpixel((x, y))
    key_pixel = get_key_pixel(state, x, y)
    outline_pixel = outline_blend(base_pixel, key_pixel)
    state.outline.putpixel((x, y), outline_pixel)


def draw_mask(state):
    # Buffer for the mask as we build, before convering to 1BPP.
    buf = state.outline.copy()
    for _ in range(state.args.radius):
        cur_buf = buf.copy()
        # Blit the image over itself in each of the 8 directions. This is equiv
        # to a [1]*9 kernel.
        directions = [(-1, -1), (-1, 1), (1, -1), (1, 1), (-1, 0), (0, -1), (0, 1), (1, 0)]
        if state.args.debug:
            directions.append((0, 0))
        for d in directions:
            buf.paste(cur_buf, d, cur_buf)
    if state.args.debug:
        mask_path = re.sub(r'(^.*)\..+$', '\\1_mask_inner.png', state.args.base)
        buf.save(mask_path)
    state.mask.paste(buf.split()[3])


def show_cv_img(img):
    cv2.imshow('image', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def find_contour(state):
    # Convert our Pillow mask image to OpenCV's format.
    img = numpy.array(state.mask.convert('L'))
    # Extract the outer contour and then approximate it.
    _, contour, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contour_length = cv2.arcLength(contour[0], True)
    epsilon = state.args.contour_factor * contour_length
    approx = cv2.approxPolyDP(contour[0], epsilon, True)
    # Create a spline interpolation.
    # Via https://stackoverflow.com/questions/14344099/smooth-spline-representation-of-an-arbitrary-contour-flength-x-y
    approx_x = approx[:,0,0]
    approx_y = approx[:,0,1]
    dist = numpy.sqrt((approx_x[:-1] - approx_x[1:])**2 + (approx_y[:-1] - approx_y[1:])**2)
    dist_along = numpy.concatenate(([0], dist.cumsum()))
    spline, u = scipy.interpolate.splprep([approx_x, approx_y], u=dist_along, s=0)
    # Resample along the spline to generate a new polygon.
    interp_d = numpy.linspace(dist_along[0], dist_along[-1], state.args.spline_factor * contour_length)
    interp_x, interp_y = scipy.interpolate.splev(interp_d, spline)
    # Convert back to a contour.
    interp_contour = numpy.array([[[int(x),int(y)]] for (x, y) in zip(interp_x, interp_y)])
    # Debugging output for the three contours.
    if state.args.debug:
        buf = cv2.cvtColor(numpy.array(state.base.convert('RGB')), cv2.COLOR_RGB2BGR)
        cv2.drawContours(buf, contour, -1, (0, 0, 255), 1)
        cv2.drawContours(buf, [approx], -1, (255, 0, 0), 1)
        cv2.drawContours(buf, [interp_contour], -1, (0, 255, 0), 1)
        contour_path = re.sub(r'(^.*)\..+$', '\\1_contour.png', state.args.base)
        cv2.imwrite(contour_path, buf)
    # Convert the interpolated contour to a bitmap.
    interp_bitmap = numpy.zeros([state.base.width, state.base.height, 1], dtype=numpy.uint8)
    cv2.drawContours(interp_bitmap, [interp_contour], -1, (255), -1)
    interp_img = Image.fromarray(interp_bitmap[:,:,0], mode='L')
    if state.args.debug:
        contour_path = re.sub(r'(^.*)\..+$', '\\1_contour_bitmap.png', state.args.base)
        interp_img.save(contour_path)
    # Create the final image.
    state.diecut = state.base.copy()
    state.diecut.putalpha(interp_img)


def init(args):
    base = Image.open(args.base).convert('RGBA')
    key = Image.open(args.key).convert('RGBA')
    return RunState(
        args=args,
        base=base,
        key=key,
        outline=Image.new('RGBA', base.size, None),
        mask=Image.new('1', base.size, 0),
        diecut=Image.new('RGBA', base.size, None),
    )


def transform(state):
    for x in range(state.base.width):
        for y in range(state.base.height):
            generate_outline(state, x, y)
    draw_mask(state)


def write(state):
    if state.args.debug:
        outline_path = re.sub(r'(^.*)\..+$', '\\1_outline.png', state.args.base)
        state.outline.save(outline_path)
        mask_path = re.sub(r'(^.*)\..+$', '\\1_mask.png', state.args.base)
        state.mask.save(mask_path)
    diecut_path = state.args.out or re.sub(r'(^.*)\..+$', '\\1_diecut.png', state.args.base)
    state.diecut.save(diecut_path)


def main(argv=sys.argv[1:]):
    parser = argparse.ArgumentParser()
    parser.add_argument('base', metavar='BASE', help='base image file')
    parser.add_argument('key', metavar='KEY', help='key image file')
    parser.add_argument('--out', '-o', metavar='FILE', help='output filename')
    parser.add_argument('--debug', '-d', help='debug mode', action='store_true')
    parser.add_argument('--radius', '-r', help='diecut blur radius', type=int, default=10)
    parser.add_argument('--contour-factor', '-c', help='contour approximation epsilon multiplier', type=float, default=0.002)
    parser.add_argument('--spline-factor', '-s', help='spline interpolation spacing mulitplier', type=float, default=1.0)
    args = parser.parse_args(args=argv)
    state = init(args)
    transform(state)
    find_contour(state)
    write(state)


if __name__ == '__main__':
    main()