Instagram Engineering Challenge: The Unshredder

unshred.rb

#!/usr/bin/env ruby -w

#------------------------------------------------------------------------------
# requires
#------------------------------------------------------------------------------

require 'RMagick'
require 'prime'

#------------------------------------------------------------------------------
# defines
#------------------------------------------------------------------------------

@kLeft  = 0
@kRight = 1

#------------------------------------------------------------------------------
# helper functions
#------------------------------------------------------------------------------

def length(color)
    # treat color as a vector, calculate the length 
    Math.sqrt(color.reduce(0) { |sum, value| sum + value * value })
end

#------------------------------------------------------------------------------

def difference(color_a, color_b)
    # convert into array, easier to manipulate
    color_a = color_a.to_hsla[0..2].map { |x| x / 256.0 }
    color_b = color_b.to_hsla[0..2].map { |x| x / 256.0 }
    # treat the color as a vector, calculate the difference
    [color_a, color_b].transpose.map { |x| x.inject :-  }
end

#------------------------------------------------------------------------------

def get_first(index, width)
    index * width
end

#------------------------------------------------------------------------------

def get_last(index, width)
    (index + 1) * width - 1
end

#------------------------------------------------------------------------------

def get_column_difference(first, second)
    [first, second].transpose.map { |x| length difference x.first, x.last }.inject :+
end

#------------------------------------------------------------------------------

#
# find_column - find the best match column 
#   side      - left or right
#   image     - image to work on
#   column    - index of column to find match for
#   available - list of available column indicies
#   width     - width of a column
#
def find_column(side, image, column, available, width)
    # configure function
    get_fst = side == @kLeft ? :get_first : :get_last
    get_lst = side == @kLeft ? :get_last  : :get_first

    # grab the first column of pixels
    first = image.get_pixels self.send(get_fst, column, width), 0, 1, image.rows

    # get difference between all other columns
    totals = available.map do |index|
        x     = self.send(get_lst, index, width)
        last  = image.get_pixels x, 0, 1, image.rows
        get_column_difference first, last
    end

    # find the least different column
    min   = totals.min
    index = totals.index(min) 
    [available[index], min]
end

#------------------------------------------------------------------------------

def find_column_width(image)
    # find all the possible multiples, since columns are equal in size
    primes, powers = image.columns.prime_division.transpose
    exponents      = powers.map{ |index| (0..index).to_a }
    divisors       = exponents.shift.product(*exponents).map do |p|
        primes.zip(p).map { |prime, power| prime ** power }.inject :*
    end

    # calculate the possible widths
    totals = divisors.sort[0..-2].map do |width|
        first  = image.get_pixels width - 1, 0, 1, image.rows
        second = image.get_pixels width, 0, 1, image.rows
        get_column_difference second, first
    end

    # find the best match, 200 seems like a valid magic number...
    totals.zip(divisors.sort).find_all { |total, divisor| total > 200.0 }[0][1]
end

#------------------------------------------------------------------------------
# main
#------------------------------------------------------------------------------

# validate arguments
if ARGV.length != 2
    puts "Usage: #{__FILE__} <shredded_image_path> <unshredded_image_path>"
    exit
end

# make sure file exist
shredded_image_path = ARGV[0]
if !File.exists? shredded_image_path then
    puts "Error: Image not found."
    exit
end

# read in the image
shredded = Magick::Image.read(shredded_image_path).first

# calculate the split value
column_width = find_column_width shredded

# calculate number of shreds
shreds = shredded.columns / column_width 

# setup variables for processing
current   = 0
available = (1..shreds-1).entries
ordered   = [current]

# loop till all the best matches have been found
lastMatch  = nil
leftMatch  = nil
rightMatch = nil
while !available.empty? 

    # if the last match was on the right calculate the right again
    if lastMatch.nil? or lastMatch == @kRight
        rightMatch = find_column @kRight, shredded, ordered[-1], available, column_width
    end

    # if the last match was on the left calculate the left again
    if lastMatch.nil? or lastMatch == @kLeft
        leftMatch = find_column @kLeft, shredded, ordered[0], available, column_width
    end

    # save the best match 
    matchedColumn = nil
    if rightMatch[1] < leftMatch[1] then
        lastMatch     = @kRight
        matchedColumn = rightMatch[0]
        ordered.push matchedColumn
    else
        lastMatch     = @kLeft
        matchedColumn = leftMatch[0]
        ordered.insert 0, matchedColumn
    end

    # remove match column from available columns
    available.delete matchedColumn

end

# create the unshredded image
unshredded = Magick::Image.new(shredded.columns, shredded.rows)
(0..shreds-1).each do |index|
    get_x  = get_first(ordered[index], column_width)
    set_x  = get_first(index, column_width)
    pixels = shredded.get_pixels get_x, 0, column_width, shredded.rows
    unshredded.store_pixels set_x, 0, column_width, shredded.rows, pixels
end

# write unshredded image
unshredded_image_path = ARGV[1]
unshredded.write unshredded_image_path