Created
January 14, 2016 10:27
-
-
Save perey/3a45d6a78cbeb62d0a09 to your computer and use it in GitHub Desktop.
Line-art painter
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python3 | |
"""Line-art painter.""" | |
# Copyright © 2016 Timothy Pederick. | |
# | |
# Permission is hereby granted, free of charge, to any person obtaining | |
# a copy of this software and associated documentation files (the | |
# "Software"), to deal in the Software without restriction, including | |
# without limitation the rights to use, copy, modify, merge, publish, | |
# distribute, sublicense, and/or sell copies of the Software, and to | |
# permit persons to whom the Software is furnished to do so, subject | |
# to the following conditions: | |
# | |
# The above copyright notice and this permission notice shall be | |
# included in all copies or substantial portions of the Software. | |
# | |
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, | |
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF | |
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. | |
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY | |
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, | |
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE | |
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. | |
# Standard library imports. | |
from argparse import ArgumentParser | |
from glob import glob | |
from pathlib import Path | |
from random import choice, shuffle | |
from string import ascii_letters, digits | |
alphanumeric = ascii_letters + digits | |
# Third-party library imports. | |
from PIL import Image, ImageDraw, ImageColor | |
# X11 colour names from Wikipedia. | |
ALL_COLOURS = frozenset( | |
{'Brown', 'RosyBrown', 'SaddleBrown', 'SandyBrown', | |
'BurlyWood', 'Chocolate', 'Peru', 'Sienna', 'Tan', | |
'Pink', 'LightPink', 'DeepPink', 'HotPink', | |
'PaleVioletRed', 'MediumVioletRed', | |
'Salmon', 'LightSalmon', 'DarkSalmon', 'Coral', 'LightCoral', 'MistyRose', | |
'Red', 'DarkRed', 'IndianRed', | |
'Crimson', 'FireBrick', 'Tomato', 'Maroon', | |
'OrangeRed', 'Orange', 'DarkOrange', | |
'Yellow', 'LightYellow', 'LightGoldenrodYellow', | |
'Goldenrod', 'PaleGoldenrod', 'DarkGoldenrod', | |
'Khaki', 'DarkKhaki', | |
'LemonChiffon', 'Moccasin', 'PapayaWhip', 'PeachPuff', | |
'YellowGreen', 'GreenYellow', | |
'Green', 'PaleGreen', 'LightGreen', 'DarkGreen', 'LimeGreen', 'Lime', | |
'SpringGreen', 'MediumSpringGreen', | |
'SeaGreen', 'LightSeaGreen', 'MediumSeaGreen', 'DarkSeaGreen', | |
'Olive', 'OliveDrab', 'DarkOliveGreen', | |
'Chartreuse', 'LawnGreen', 'ForestGreen', | |
'Aquamarine', 'MediumAquamarine', | |
'Cyan', 'LightCyan', 'DarkCyan', 'Teal', | |
'Turquoise', 'PaleTurquoise', 'MediumTurquoise', 'DarkTurquoise', | |
'Blue', 'LightBlue', 'MediumBlue', 'DarkBlue', | |
'SkyBlue', 'LightSkyBlue', 'DeepSkyBlue', | |
'SlateBlue', 'MediumSlateBlue', 'DarkSlateBlue', | |
'SteelBlue', 'LightSteelBlue', | |
'AliceBlue', 'CadetBlue', 'CornflowerBlue', 'DodgerBlue', 'MidnightBlue', | |
'PowderBlue', 'RoyalBlue', 'Navy', 'Indigo', | |
'BlueViolet', 'Violet', 'DarkViolet', 'Magenta', 'DarkMagenta', | |
'Purple', 'MediumPurple', 'Orchid', 'MediumOrchid', 'DarkOrchid', | |
'Lavender', 'LavenderBlush', 'Plum', 'Thistle', | |
'White', 'AntiqueWhite', 'FloralWhite', 'GhostWhite', 'NavajoWhite', | |
'Azure', 'Beige', 'Bisque', 'BlanchedAlmond', 'Cornsilk', 'Honeydew', | |
'Ivory', 'MintCream', 'Linen', 'OldLace', 'Seashell', 'Snow', 'Wheat', | |
'Gray', 'LightGray', 'DimGray', 'DarkGray', | |
'SlateGray', 'LightSlateGray', 'DarkSlateGray', | |
'Gainsboro', 'Silver', 'WhiteSmoke', 'Black' | |
}) | |
SKY_COLOURS = frozenset( | |
{'Aquamarine', 'Cyan', 'LightCyan', 'PaleTurquoise', 'LightBlue', | |
'SkyBlue', 'LightSkyBlue', 'DeepSkyBlue', | |
'SlateBlue', 'SteelBlue', 'LightSteelBlue', | |
'AliceBlue', 'CornflowerBlue', 'DodgerBlue', | |
'PowderBlue', 'RoyalBlue' | |
}) | |
VEGE_COLOURS = frozenset( | |
{'YellowGreen', 'GreenYellow', 'Green', 'LightGreen', 'DarkGreen', | |
'LimeGreen', 'SpringGreen', 'MediumSpringGreen', | |
'SeaGreen', 'DarkSeaGreen', 'OliveDrab', 'DarkOliveGreen', | |
'LawnGreen', 'ForestGreen', 'MediumAquamarine' | |
}) | |
EARTH_COLOURS = frozenset( | |
{'Brown', 'RosyBrown', 'SaddleBrown', 'SandyBrown', | |
'BurlyWood', 'Chocolate', 'Peru', 'Sienna', 'Tan' | |
}) | |
RESERVED_COLOURS = frozenset( | |
{'White', 'Black' | |
}) | |
OTHER_COLOURS = (ALL_COLOURS - SKY_COLOURS - VEGE_COLOURS - | |
EARTH_COLOURS - RESERVED_COLOURS) | |
# Greyscale and alpha-level keywords. | |
BLACK = lambda mode: {'RGB': (0, 0, 0), 'RGBA': (0, 0, 0, 255), | |
'LA': (0, 255)}.get(mode, 0) | |
WHITE = lambda mode: {'RGB': (255, 255, 255), 'RGBA': (255, 255, 255, 255), | |
'LA': (255, 255)}.get(mode, 255) | |
CLEAR = lambda mode: {'RGBA': (255, 255, 255, 0), | |
'LA': (255, 0)}.get(mode, WHITE(mode)) | |
# Utility functions. | |
def paste_with_alpha(im1, im2): | |
"""Paste im2 into im1, using im2's alpha channel. | |
This is an in-place operation modifying im1, like Pillow's | |
Image.paste(), but unlike that operation the alpha channel of im2 | |
is not ignored. It is, however, presumed to exist. | |
""" | |
*_, alpha_layer = im2.split() | |
im1.paste(im2, mask=alpha_layer) | |
# For distance minimisation purposes, we don't need to take the square root of | |
# the sum of squares, so don't bother. | |
distance = lambda p1, p2: sum((a - b) ** 2 for a, b in zip(p1, p2)) | |
def nearest_other(img, replace_colour=None, search_dist=5): | |
"""Replace pixels of a given colour with the nearest other colour. | |
The image is not modified in-place; a new RGB image is returned with | |
the pixels replaced. Note that the alpha channel of the original | |
image is not included. | |
""" | |
# If no colour is specified, use black. | |
if replace_colour is None: | |
replace_colour = BLACK(img.mode) | |
# Record the distances to unreplaced pixels using the alpha channel. To | |
# initialise the alpha channel to maximum, remove and re-add it. | |
recoloured = img.convert('RGB').convert('RGBA') | |
for x in range(img.width): | |
for y in range(img.height): | |
# If this pixel has an unreplaced colour, mark it as the nearest | |
# neighbour to nearby pixels... unless they already have a shorter | |
# distance. | |
colour = img.getpixel((x, y)) | |
if colour != replace_colour: | |
r, g, b, _ = recoloured.getpixel((x, y)) | |
recoloured.putpixel((x, y), (r, g, b, 0)) | |
for sx in range(max(0, x - search_dist), | |
min(img.width, x + search_dist + 1)): | |
for sy in range(max(0, y - search_dist), | |
min(img.height, y + search_dist + 1)): | |
if (sx, sy) == (x, y): | |
continue | |
if img.getpixel((sx, sy)) == replace_colour: | |
*_, old_dist = recoloured.getpixel((sx, sy)) | |
new_dist = distance((x, y), (sx, sy)) | |
if new_dist < old_dist: | |
recoloured.putpixel((sx, sy), | |
(r, g, b, new_dist)) | |
# Replace all of the working-out values in the alpha channel with full | |
# opacity. | |
bilevel_alpha = lambda alpha: 255 if alpha else 0 | |
recoloured.putalpha(recoloured.split()[-1].point(bilevel_alpha)) | |
return recoloured | |
_random_colours_used = set(RESERVED_COLOURS) | |
def random_colour(position, image): | |
"""Choose a random colour for a region of an image. | |
To ensure that a range of colours is used, no colour will be | |
returned more than once, unless and until a set of colours has been | |
exhausted. | |
Keyword arguments: | |
position -- A 2-tuple giving the x and y coordinates of the | |
location in the image for which a colour needs assigning. | |
image -- The image to be coloured. | |
""" | |
global _random_colours_used | |
# Are there any colours still available? | |
if _random_colours_used >= ALL_COLOURS: | |
# No. Make all non-reserved colours available again. | |
_random_colours_used = set(RESERVED_COLOURS) | |
# Pick a colour at random. | |
chosen_colour = choice(tuple(ALL_COLOURS - _random_colours_used)) | |
_random_colours_used.add(chosen_colour) | |
# Return the colour's value. | |
return ImageColor.getrgb(chosen_colour) | |
_natural_colours_used = set(RESERVED_COLOURS) | |
def natural_colour(position, image): | |
"""Guess an appropriate natural colour for a region of an image. | |
The image is divided into thirds vertically. A variety of colours is | |
used for the middle third. In the outer thirds, sky colours are used | |
in the top ninths, earth colours in the middle ninths, and grass | |
colours in the bottom ninths. | |
Additionally, to ensure that a range of colours is used, no colour | |
will be returned more than once, unless and until a set of colours | |
has been exhausted. | |
Keyword arguments: | |
position -- A 2-tuple giving the x and y coordinates of the | |
location in the image for which a colour needs assigning. | |
image -- The image to be coloured. | |
""" | |
global _natural_colours_used | |
# Which set of colours should be chosen from? | |
sectors = [[SKY_COLOURS, OTHER_COLOURS, SKY_COLOURS], | |
[EARTH_COLOURS, OTHER_COLOURS, EARTH_COLOURS], | |
[VEGE_COLOURS, OTHER_COLOURS, VEGE_COLOURS]] | |
x, y = position | |
width, height = image.size | |
colours = sectors[3 * y // height][3 * x // width] | |
# Are there any colours still available in this set? | |
if _natural_colours_used >= colours: | |
# No. Make this whole set of colours available again. | |
_natural_colours_used -= colours | |
# Pick a colour at random. | |
chosen_colour = choice(tuple(colours - _natural_colours_used)) | |
_natural_colours_used.add(chosen_colour) | |
# Return the colour's value. | |
return ImageColor.getrgb(chosen_colour) | |
def paint(image, palette_name, threshold=0.5, proper_fill=False, | |
save_debug=False): | |
"""Paint a line-art image. | |
Keyword arguments: | |
image -- The image to be painted. | |
threshold -- The brightness cutoff between the lines and the | |
paintable areas of the image, expressed as a proportion | |
between 0 and 1 between the darkest and lightest colours in | |
the image. The default is 0.5 (i.e. halfway between the | |
extremes). | |
save_debug -- Whether or not to save intermediate results (in | |
the current working directory). The default is False. | |
""" | |
if save_debug: | |
# Generate a 6-character random alphanumeric prefix. | |
random_prefix = lambda length: ''.join(choice(alphanumeric) | |
for _ in range(length)) | |
prefix = random_prefix(6) | |
# Ensure the prefix isn't currently used. | |
while glob('{}*.png'.format(prefix)): | |
prefix = random_prefix(6) | |
# Convert the input image to greyscale, no alpha. First, does it have an | |
# alpha channel to dispose of? | |
if 'A' in image.getbands(): | |
# Yes; compose it onto a white background. | |
all_white = Image.new(image.mode, image.size, WHITE(image.mode)) | |
paste_with_alpha(all_white, image) | |
image = all_white | |
# Convert the colour mode. | |
image = image.convert('L') | |
# Identify the darkest and lightest colours in the input image. | |
darkest, lightest = image.getextrema() | |
# Create a black-and-white copy of the input image, with the cutoff level | |
# at some percentage lighter than the darkest colour. | |
cutoff = int((1 - threshold) * darkest + threshold * WHITE(image.mode)) | |
bilevel = lambda colour: (BLACK if colour < cutoff else WHITE)(image.mode) | |
paintable = image.point(bilevel) | |
if save_debug: | |
paintable.save('{}-bilevel.png'.format(prefix)) | |
# Convert the bi-level image to full-colour so that it can be painted. | |
paintable = paintable.convert('RGB') | |
# Find and fill all white regions. | |
colour_fn = {'natural': natural_colour}.get(palette_name, random_colour) | |
white = WHITE(paintable.mode) | |
for x in range(paintable.width): | |
for y in range(paintable.height): | |
if paintable.getpixel((x, y)) == white: | |
ImageDraw.floodfill(paintable, (x, y), | |
colour_fn((x, y), paintable)) | |
if save_debug: | |
paintable.save('{}-painted.png'.format(prefix)) | |
# Fill black pixels with their nearest non-black neighbour. | |
if proper_fill: | |
blackfill = nearest_other(paintable) | |
else: | |
blackfill = paintable.convert('RGBA') | |
black, clear = BLACK(blackfill.mode), CLEAR(blackfill.mode) | |
while any(colour == black for count, colour in blackfill.getcolors()): | |
for x in range(blackfill.width): | |
for y in range(blackfill.height): | |
if paintable.getpixel((x, y)) == black: | |
# Check neighbouring pixels in a random order, to try | |
# and avoid too much colour bleed from any one side. | |
neighbours = [(x + 1, y), (x, y + 1), | |
(x - 1, y), (x, y - 1)] | |
shuffle(neighbours) | |
for neighbour_pos in neighbours: | |
try: | |
neighbour = paintable.getpixel(neighbour_pos) | |
except IndexError: | |
pass | |
else: | |
if neighbour != black: | |
blackfill.putpixel((x, y), neighbour) | |
break | |
else: | |
# Fill all other pixels with fully transparent white. | |
blackfill.putpixel((x, y), clear) | |
paste_with_alpha(paintable, blackfill) | |
if save_debug: | |
blackfill.save('{}-blackfill.png'.format(prefix)) | |
paintable.save('{}-allpainted.png'.format(prefix)) | |
# Create an alpha mask from the input image: completely opaque where the | |
# colour is darkest, fading to fully transparent as it pales. | |
fade = lambda colour: WHITE(image.mode) * ((colour - lightest) / | |
(darkest - lightest)) | |
mask = image.point(fade) | |
if save_debug: | |
mask.save('{}-mask.png'.format(prefix)) | |
# Compose the input image (using the alpha mask) over the coloured image, | |
# returning the result. | |
return Image.composite(image, paintable, mask).convert('RGB') | |
def parse_args(): | |
"""Parse command-line arguments.""" | |
parser = ArgumentParser(description='Paint one or more line-art images.') | |
parser.add_argument('-p', '--palette', default='random', help='a palette ' | |
'from which to choose colours; one of "random" (the ' | |
'default) or "natural"') | |
parser.add_argument('-t', '--threshold', type=float, default=0.5, | |
help='the lightness threshold between outlines and ' | |
'paintable areas (a proportion from 0 to 1)') | |
fill_mode = parser.add_mutually_exclusive_group() | |
fill_mode.add_argument('-f', '--proper-fill', action='store_true', | |
help='fill under black lines with proper nearest-' | |
'neighbour searching (slow)') | |
fill_mode.add_argument('-F', '---no-proper-fill', action='store_false', | |
dest='proper_fill', help='fill under black lines ' | |
'with approximate nearest-neighbour searching ' | |
'(fast)') | |
parser.add_argument('--format', help='the image format to use for the ' | |
'result (the default is the same as the input file)') | |
parser.add_argument('-d', '--debug', action='store_true', help='output ' | |
'debugging information') | |
parser.add_argument('images', metavar='FILE', nargs='+', | |
help='one or more image filenames') | |
return parser.parse_args() | |
def main(): | |
"""Run the command-line program.""" | |
args = parse_args() | |
for filename in args.images: | |
# Get a new filename for the result: name.ext → name.painted.ext | |
filepath = Path(filename) | |
name_part, extensions = filepath.stem, filepath.suffixes | |
if args.format is not None: | |
extensions[-1] = '.' + args.format | |
new_filename = ''.join([name_part, '.painted'] + extensions) | |
# Open the original file. | |
original = Image.open(filename) | |
try: | |
painted = paint(original, args.palette, args.threshold, | |
args.proper_fill, args.debug) | |
painted.save(new_filename) | |
finally: | |
original.close() | |
if __name__ == '__main__': | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment