Instagram Engineering Challenge: The Unshredder

instagram-unshredder.py

#find out more http://instagram-engineering.tumblr.com/post/12651721845/instagram-engineering-challenge-the-unshredder
from PIL import Image
from collections import defaultdict
import operator

def unshred(file_name, NUMBER_OF_COLUMNS, diff, num_check_points=0): 
    """
    @params diff base on this value we decide if two pixels are beside each other or not
    @params num_check_points how many times we check pixels across shreds
    """
    image = Image.open(file_name)
    if num_check_points == 0 :
        num_check_points = image.size[1]//40
    def log(*data):
        #print data
        pass

    # Create a new image of the same size as the original
    # and copy a region into the new image
    unshredded = Image.new("RGBA", image.size)
    shred_width = unshredded.size[0]/NUMBER_OF_COLUMNS

    #keep shreds 
    shreds = []
    for i in xrange(NUMBER_OF_COLUMNS):
        shred_number = i
        x1, y1 = shred_width * shred_number, 0
        x2, y2 = x1 + shred_width, image.size[1]
        shreds.append( image.crop((x1, y1, x2, y2)) )

    def find_peace(current, which):
        """find next or previous shreds
        We need both to estimate starting point of photo

        We estimate possibility of side shreds based on @diff and @num_check_points, the shred with highest possibility would be our answer, it's rough guess but apparently it works!
        
        more smart algorithm could check all shreds possibility and decide which one is worng

        @return next or previous shreds sorted by their possibility
        """
        a = defaultdict(lambda : 0)
        if which == 'RIGHT':
            x_current = shred_width-1
            x_beside   = 0 
        elif which == "LEFT" :
            x_current = 0
            x_beside   = shred_width-1
        for y in xrange(0,image.size[1], num_check_points):
            x_y   = shreds[current].getpixel((x_current, y))
            for i in xrange(0, NUMBER_OF_COLUMNS):
                #don't check points on current shred
                if current == i :
                    continue
                o_x_y = shreds[i].getpixel((x_beside ,y))
                if abs(x_y[0]-o_x_y[0]) < diff and abs(x_y[1]-o_x_y[1]) < diff and abs(x_y[2]-o_x_y[2]) < diff and abs(x_y[3]-o_x_y[3]) < diff :
                    a[i]+=1

        return sorted(a.iteritems(), key=operator.itemgetter(1))


    left_sides = []
    for i in xrange(0, NUMBER_OF_COLUMNS):
        pr = find_peace(i, which='LEFT')[-1:]
        log("For chunk", i , pr )
        left_sides.append(pr[0])
    log(left_sides)

    right_sides = []
    for i in xrange(0, NUMBER_OF_COLUMNS):
        pr = find_peace(i, which='RIGHT')[-1:]
        log("For chunk", i , pr)
        right_sides.append(pr[0])
    log(right_sides)

    def get_first_shreds(num_shreds, right_match, left_match):
        item = 0
        for i in xrange(num_shreds):
            right_shreds = right_match[i]
            left_shreds  = left_match[right_shreds[0]]
            if not i == left_shreds[0]:
                item = right_shreds[0]
        return item 

    item = first_item = get_first_shreds(NUMBER_OF_COLUMNS, right_sides, left_sides)
    already_added = set()
    i, valve, current_width = [0] * 3
    while i < NUMBER_OF_COLUMNS :
        valve+=1
        if not item in already_added  :
            destination_point = (current_width, 0)
            unshredded.paste(shreds[item], destination_point)
            already_added.add(item)
            current_width+=shred_width
            item = right_sides[item][0]
            i+=1
        elif valve> NUMBER_OF_COLUMNS *2:
            print "Oops! it might not be the full image but it's close\nTry to play with diff and num_check_points args and find your match!"
            break

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

unshred('TokyoPanoramaShredded.png', 20, diff=40)