Take a random image and replace all distinct colours with a random other colour

colorise.py

"""Take a random image and replace all distinct colours with
a random other colour

:arg infile:  The source image to read from
:arg loop:    Whether to loop the list of colours instead of only replacing the N first
:arg upper:   The upper bound of colours to replace
"""

from PIL import Image
from collections import defaultdict
import sys

# Parse arguments
argv = list(sys.argv)
infile = argv[1] if 1 < len(argv) else "image.jpg"
loop = bool(argv[2]) if 2 < len(argv) else False
INF = int(argv[3]) if 3 < len(argv) else None
if INF is not None and INF < 0:
    INF = float("inf")

# Build list of colours
colors = []
z = {128, 0, 255, 64, 192}
for r in z:
    for g in z - {r}:
        for b in z - {r, g}:
            c = (r, g, b)
            colors.append(c)

# Open input image
# Please note that this is a O(n^2 + n^3) operation (ie. slow)
with Image.open(infile) as im:
    print(f"To process: {im.width * im.height} pixels")

    # Count all distinct pixels
    # TODO: use collections.Counter instead??
    pixels = defaultdict(int)
    for y in range(0, im.height):
        for x in range(0, im.width):
            pi = im.getpixel((x, y))
            pixels[pi] += 1
    print(f"{len(pixels)} distinct colours found")

    # Sort based on number
    s = sorted(pixels.items(), key=lambda x: x[1], reverse=True)

    # Calculate upper bound
    if INF is None:
        # If none given, replace half of all available colours
        # Note that this is not the same as half of all pixels.
        #   For most natural JPEG-compressed images, this seems
        #   to replace about 95% of all pixels.
        INF = len(s) // 2
    upper = min(len(s), len(colors) if not loop else INF)
    print(f"Building the replacement map for {upper} colours")
 
    # Build the replacement map
    mapp = dict()
    for i in range(0, upper):
        mapp[s[i][0]] = colors[i % len(colors)]

    # Apply the replacements
    # XXX: This is an O(n^3) loop and I'm not sure if we can do better
    print(f"Replacing all pixels")
    for y in range(0, im.height):
        for x in range(0, im.width):
            xy = (x, y)
            pi = im.getpixel(xy)
            if pi in mapp:
                im.putpixel(xy, mapp[pi])

    # Write to output file
    outfile = infile[::-1].replace(".", ".mod."[::-1])[::-1]
    print(f"Writing to {outfile}")
    im.save(outfile)