Instagram Coding Challenge Unshredder

unshred.py

from PIL import Image
from math import sqrt
import sys

MASK_SIZE = 3
THRESHOLD = 0.9

def get_shred_width():
    width = image.size[0]

    # Divide up the image into pixel strips and calculate the difference between each adjacent strip
    strips = [get_pixel_strip(x, mask_size=1) for x in xrange(width)]
    differences = [distance(strips[i], strips[i+1]) for i in xrange(width - 1)]
    
    # Normalize the difference values by computing the number of standard deviations away from 
    # the mean for each of the calculated differences.  If the normalized difference
    # exceeds a threshold value then it is considered a valid shred boundary.
    mean = sum(differences) / len(differences)
    std_dev = sqrt(sum([pow(x - mean, 2) for x in differences]) / len(differences))
    normalized_diffs = [abs(x - mean) / std_dev for x in differences]
    
    for shred_width in xrange(1, width):
        # Only consider columns that are evenly distributed and uniformly sized
        if width % shred_width == 0:
            boundary_diffs = [normalized_diffs[i-1] for i in xrange(shred_width, width, shred_width)]
            if all([d >= THRESHOLD for d in boundary_diffs]):
                return shred_width
    
    # Couldn't determine the shred width, use the default.
    return 32
        
def color_average(pixels):
    """ Returns the average color for the list of (r,g,b,a) tuples passed in."""
    size = float(len(pixels))
    return tuple([sum(x) / size for x in zip(*pixels)])
    
def color_distance(c, c2):
    """ Calculates the squared Euclidean distance between the two specified
    (r, g, b, a) color tuples """ 
    return sum([pow(x - y, 2) for x, y in zip(c, c2)])
    
def distance(s, s2):
    """ Returns the distance between two pixel strips """
    return sum([color_distance(c, c2) for c, c2 in zip(s, s2)])

def get_pixel_value(x, y):
    width = image.size[0]
    pixel = data[y * width + x]
    return pixel

def get_pixel_strip(x, mask_size=MASK_SIZE):
    height = image.size[1]
    strip = [] 
    for y in xrange(0, height - mask_size, mask_size):
        pixel_mask = []
        for xoffset in xrange(mask_size):
            for yoffset in xrange(mask_size):
                pixel_mask.append(get_pixel_value(x + xoffset, y + yoffset))
        strip.append(color_average(pixel_mask))
        
    return strip

image = Image.open(sys.argv[1])
data = image.getdata()

shred_width = get_shred_width()

# Each shred is a tuple consisting of (left pixel strip, right pixel trip)
shreds = [(get_pixel_strip(x), 
           get_pixel_strip(x + shred_width - 1 - MASK_SIZE)) for x in xrange(0, image.size[0], shred_width)]

first_shred = None
first_shred_dist = None

shred_distances = []
for i, shred in enumerate(shreds):
    # Calculate the distance between the left edge of this shred and the right 
    # edge of all other shreds.  Generate a list of (shred_index, next_shred_index, distance) tuples
    matches = [(j, i, distance(shred[0], s[1])) for j, s in enumerate(shreds) if s != shred]

    shred_distances.extend(matches)
                
    # Check to see if i is the first shred in the image.
    # The first shred in the image will not have a previous image so it will
    # be "far away" from all of the other shreds.
    min_dist = min([m[2] for m in matches])

    if not first_shred or min_dist > first_shred_dist:
        first_shred = i
        first_shred_dist = min_dist

shred_distances.sort(key=lambda x: x[2])

# Greedily pick out the next shred 
next_shred = {}
for p, n, dist in shred_distances:
    if not p in next_shred and not n in next_shred.values():
        next_shred[p] = n
        
# Write out the destination image.
unshredded = Image.new("RGBA", image.size)

curr_shred = first_shred
for x in xrange(0, image.size[0], shred_width):
    x1, y1 = shred_width * curr_shred, 0
    x2, y2 = x1 + shred_width, image.size[1]
    source_region = image.crop((x1, y1, x2, y2))
    destination_point = (x,0)
    unshredded.paste(source_region, destination_point)

    # Find the next shred   
    curr_shred = next_shred[curr_shred] 

unshredded.save("unshredded.png", "PNG")