Convert an image to all-RGB by colourizing from 3d hilbert line

allrgbify.py

import Image    # Python Imaging Library (PIL)
import random
import math
import os
import sys
import getopt


# Main work function. Takes the path to a (square) image and the number of
# bits per colour channel desired. That value dictates the size of the output
# image, so it must be divisible by 2. For < 8, lowest order bits dropped.
def allrgbify(path, bitsPerChannel):
    if bitsPerChannel & 1 != 0:
        print "bitsPerChannel must be divisible by 2"
        return
    if (bitsPerChannel > 8):
        print "Do you really want a picture " + \
            str(int(math.sqrt(2**(3*bitsPerChannel)))) + " pixels on a side?"
        return
    size = 2**(3*bitsPerChannel/2)
    shift = 8-bitsPerChannel

    im = Image.open(path)
    im = im.resize((size, size))
    pix = im.load()

    # Flat is an array of tuples, (index, x, y)
    flat = [(hindex(bitsPerChannel, map(lambda a: a>>shift, pix[x,y])), x, y)
                for y in xrange(size) for x in xrange(size)]

    # Shuffle so pixels with the same calculated index treated "fairly".
    random.shuffle(flat)

    # Sort by calculated index function.
    def cmp(x,y):
        return x[0]-y[0]
    flat.sort(cmp)

    # Map new colour scheme back to the respective pixels in the image.
    for i in xrange(2**(3*(8-shift))):
        pt = flat[i]
        pix[pt[1], pt[2]] = tuple(map(lambda a: a<<shift,
                                  hindexInverse(bitsPerChannel, i)))

    name, ext = os.path.splitext(path)
    im.save(name + "_allrgbify.png", "PNG")

# Generate allrgb versions of selected images with specified bit depth.
def genAll():
    files = ["C:/Users/Eric Burnett/Desktop/Python/zoom5.png",
             "C:/Users/Eric Burnett/Desktop/Python/zoom6.png"]

    for f in files:
        allrgbify(f, 6)

# Semi-related function - output the hilbert colours in a line.
def hilbertLine():
    im = Image.new("RGB", (512, 1))
    im.putdata([tuple(map(lambda a: a<<5, hindexInverse(3, i)))
                    for i in range(512)])
    im.save("C:/Users/Eric Burnett/Desktop/Python/hilbertLine.png")

# Verify a file to ensure it actually uses every colour only once. Only works
# for images generated with bit depth 8.
def verify(path):
    im = Image.open(path)

    pix = im.load()
    size = im.size
    flat = [(pix[x,y][0]<<16)|(pix[x,y][1]<<8)|pix[x,y][2]
                for y in range(size[0]) for x in range(size[1])]
    flat.sort()
    for i in xrange(2**24-1):
        if flat[i] != i:
            print "incorrect value! i="+str(i)
            return

# Hilbert logic from https://www.cs.dal.ca/sites/default/files/CS-2006-07.pdf
# Algorithm a direct implementation of the logic presented there.
# returns the index into the rgb cube indexed by the 3d hilbert curve
def hindex(m, p):
    dim = 3
    h = 0
    e = 0
    d = 0
    for i in range(m)[::-1]:
        l = bit(p[2], i) << 2 | bit(p[1], i) << 1 | bit(p[0], i)
        l = cycleLeft(l, 1, dim)
        l = T(e, d, dim, l)
        w = gcInverse(dim, l)
        e = e ^ (cycleLeft(E(w), d+1, dim))
        d = (d + D(w, dim) + 1) % dim
        h = (h << dim) | w
    return h

def hindexInverse(m, h):
    dim = 3
    e = 0
    d = 0
    p = [0, 0, 0]
    for i in range(m)[::-1]:
        w = bit(h, 3*i+2) << 2 | bit(h, 3*i+1) << 1 | bit(h, 3*i)
        l = gc(w)
        l = TInverse(e, d, dim, l)
        for j in range(dim):
            p[j] |= bit(l, j) << i
        e = e ^ (cycleLeft(E(w), d+1, dim))
        d = (d + D(w, dim) + 1) % dim
    return p

# trailing set bits
def tsb(i):
    # Split on last 4 bits
    if ((i & 7) == 7):
        # All 4 set
        return 4 + tsb(i >> 4)
    tsb4 = (0, 1, 0, 2, 0, 1, 0, 3)
    return tsb4[i & 7]

# the greycode of the integer i
def gc(i):
    return i ^ (i >> 1)

def gcInverse(m, g):
    i = g
    j = 1
    while j < m:
        i = i ^ (g >> j)
        j = j + 1
    return i

def D(i, n):
    if i == 0:
        return 0
    if i & 1:
        return tsb(i) % n
    return tsb(i-1) % n

def E(i):
    if i == 0:
        return 0
    return gc((i-1) & (~1))

def cycle(a, b, n):
    return (a >> b) | ((a << (n-b)) & (2**n - 1))

def cycleLeft(a, b, n):
    return cycle(a, n-b, n)

def T(e, d, n, b):
    return cycle(b^e,d+1,n)

def TInverse(e, d, n, b):
    return T(cycle(e, d+1, n), n-d-1, n, b)

def bit(n, i):
    return (n >> i) & 1

def bitSet(n, i, val):
    return n | val << i

# Ensure hilbert mapping is a bijection as expected.
def testHilbertInvertible():
    for i in xrange(2**24-1):
        if hindex(8, hindexInverse(8, i)) != i:
            print "failure on ", i
            return 0
            
def main():
    possibleArgs = ["image=", "bpc="]
    usage = "Usage: python allrgbify.py --image=\"image path\" --bpc=desired_bits_per_channel"
    
    try:
        opts, args = getopt.getopt(sys.argv[1:], "", possibleArgs)
    except getopt.error, msg:
        print msg
        print usage
        sys.exit(2)
        
    path = ""
    bpc = 0
    for o, a in opts:
        if o == "--image":
            path = a
        else:
            bpc = int(a)
            
    if path == "":
		print usage
		sys.exit(2)
    allrgbify(path, bpc)

if __name__ == "__main__":
    main()