Instagram unshredder challenge

unshred.py

#!/usr/bin/python2
# unshred.py - Vedant Kumar <[email protected]>

import sys
import math
import colorsys
from collections import deque

from PIL import Image

def deshred(fin, fout):
	'Reads a shredded image from fin, writes the original to fout.'
	img = Image.open(fin)
	buffer = img.getdata()
	hls_buffer = [colorsys.rgb_to_hls(*[cell / 255.0 for cell in rgba[:3]])
		for rgba in buffer]
	width, height = img.size

	def getpixel(col, row):
		'Gets an HLS 3-tuple.'
		return hls_buffer[row * width + col]

	def pix_diff(lhp, rhp):
		'Finds the difference between two pixels.'
		return math.sqrt(sum([
			(lhp[k] - rhp[k]) ** 2 for k in xrange(len(lhp))
		]))

	def col_diff(lhc, rhc):
		'Finds the badness-of-fit between two columns.'
		return sum([pix_diff(
			getpixel(lhc, row), getpixel(rhc, row)
		) for row in xrange(height)])

	def shred_widths():
		'Finds the width of the shreds in the image.'
		print "Finding shred width..."
		deltas = [col_diff(x, x + 1) ** 3 for x in xrange(width - 1)]

		def is_max_delta(j):
			'Checks if j is a local maximum in the deltas function.'
			return deltas[j] > deltas[j - 1] and deltas[j] > deltas[j + 1]

		def next_max_delta(j, factor):
			'Finds the next potential max in deltas.'
			return (j + 1) * factor - 1

		for i in xrange(1, len(deltas) - 1):
			if is_max_delta(i):
				factor = 2
				goodMaximum = True
				while next_max_delta(i, factor) < len(deltas):
					new = next_max_delta(i, factor)
					if not is_max_delta(new):
						goodMaximum = False
						break
					factor += 1
				if goodMaximum:
					w = i + 1
					print "Determined shred width:", w
					return w
		return 1

	shred_width = shred_widths()
	nr_shreds = width / shred_width

	def left_column(shred):
		return shred * shred_width

	def right_column(shred):
		return left_column(shred + 1) - 1

	posns = deque([0])
	unvisited = range(1, nr_shreds)
	while len(unvisited):
		lhs, lscore = min([
			(lhs, col_diff(left_column(posns[0]), right_column(lhs)))
				for lhs in unvisited], key=lambda pair: pair[1])
		rhs, rscore = min([
			(rhs, col_diff(right_column(posns[-1]), left_column(rhs)))
				for rhs in unvisited], key=lambda pair: pair[1])
		if lscore < rscore:
			posns.appendleft(lhs)
			unvisited.remove(lhs)
		else:
			posns.append(rhs)
			unvisited.remove(rhs)
		print "posns =", posns

	orig = Image.new("RGBA", img.size)
	for new, shred in enumerate(posns):
		x1, y1 = shred * shred_width, 0
		x2, y2 = x1 + shred_width, height
		orig.paste(img.crop((x1, y1, x2, y2)), (new * shred_width, 0))
	orig.save(fout, "PNG")

if __name__ == '__main__':
	if len(sys.argv) == 3:
		deshred(*sys.argv[1:])
	else:
		deshred("in.png", "out.png")