Image Rearrannger

import os
import sys
import Image

#returns dict of (color -> pixel list) pairs
def imgColorLocs(img):
    data = img.load()
    colorLocs = {}
    for x in range(img.size[0]):
        for y in range(img.size[1]):
            color = data[x, y]
            if color in colorLocs:
                colorLocs[color].append((x, y))
            else:
                colorLocs[color] = [(x, y)]
    return colorLocs

#returns a tuple of 3 ordered pairs: the canvas size, the src position, and the dst position
def getCanvasDims(src, dst, inPlace):
    w = max(src.size[0], dst.size[0])
    h = max(src.size[1], dst.size[1])
    offset = lambda dim, canvasDim: int(0.5 * (canvasDim - dim))
    srcPos = offset(src.size[0], w), offset(src.size[1], h)
    dstPos = offset(dst.size[0], w), offset(dst.size[1], h)
    if not inPlace:
        dstPos = dstPos[0] + w, dstPos[1]
        w *= 2
    return (w, h), srcPos, dstPos

def dist(loc1, loc2):
    return ((loc2[0] - loc1[0])**2 + (loc2[1] - loc1[1])**2)**0.5

#returns list of tuples of the closest color ((srcx, srcy), (dstx, dsty), (r, g, b))
def getMappings(src, srcOffset, dst, dstOffset):
    add = lambda t1, t2: (t1[0] + t2[0], t1[1] + t2[1])
    srcColorLocs = imgColorLocs(src)
    dstColorLocs = imgColorLocs(dst)
    mappings = []
    for color in srcColorLocs:
        srcLocs = srcColorLocs[color]
        dstLocs = dstColorLocs[color]
        for srcLoc in srcLocs:
            minDist = dist((0, 0), src.size)
            index = 0
            for i, dstLoc in enumerate(dstLocs):
                d = dist(srcLoc, dstLoc)
                if d < minDist:
                    minDist = d
                    index = i
            dstLoc = dstLocs.pop(index)
            mappings.append((add(srcLoc, srcOffset), add(dstLoc, dstOffset), color))
    return mappings

def makeLinearMapper(steps):
    def linearMapper(startLoc, stopLoc, step):
        m = lambda start, stop: (stop - start) * step / steps + start
        return (m(startLoc[0], stopLoc[0]), m(startLoc[1], stopLoc[1]))
    return linearMapper

#src is the filename of the palette image that will be rearranged
#dst if the filename of the final rearranged palette
#steps is the number of increments between src and dst, steps + 1 images are generated in total
#folder is the location all intermediate images will save to
#bg is the background of the intermediate images
#when cycle is true a duplicate set of images will be made with higher numbers in reverse order, useful for making continuous looking gifs
#when inPlace is true the transforming images are drawn on top of one another
#when drawSrc is true the src image is drawn below the moving pixels
#when drawDst is true the dst image is drawn below the moving pixels
def go(src, dst, steps, folder='output', bg='black', cycle=False, inPlace=True, drawSrc=False, drawDst=False):
    print 'STARTING'

    def saveImg(img, number):
        img.save('%s/img_%04d.png' % (folder, number))

    src = Image.open(src).convert('RGB')
    dst = Image.open(dst).convert('RGB')
    if sorted(src.getdata()) != sorted(dst.getdata()):
        print 'INCOMPATIBLE IMAGES'
        sys.exit()

    if not os.path.exists(folder):
        os.makedirs(folder)

    canvasDims = getCanvasDims(src, dst, inPlace)
    mappings = getMappings(src, canvasDims[1], dst, canvasDims[2])
    mapper = makeLinearMapper(steps)

    for step in range(steps + 1):
        print 'Generating image %d of %d' % (step + 1, steps + 1)

        img = Image.new('RGB', canvasDims[0], bg)
        if drawSrc:
            img.paste(src, canvasDims[1])
        if drawDst:
            img.paste(dst, canvasDims[2])
        data = img.load()
        for startLoc, stopLoc, color in mappings:
            x, y = mapper(startLoc, stopLoc, step)
            data[x, y] = color
        saveImg(img, step)

        if cycle and step != 0 and step != steps:
            saveImg(img, 2 * steps - step)

    print 'DONE'

go('src.png', 'dst.png', 19)