image unshredder for instagram engineering challenge

unshredder.rb

#!/usr/bin/env ruby
#
# image unshredder for instagram engineering challenge:
#   <http://instagram-engineering.tumblr.com/post/12651721845>
#
# author: @mikedemers <[email protected]>
#
# t-shirt size: large :)

require 'rubygems'
require 'RMagick'

raise "Usage: #{$0} INPUT_IMAGE OUTPUT_IMAGE" if ARGV.size < 2

class Slice
  attr_reader :image, :width, :height
  
  def left_edge_image
    @left_edge_image ||= image.crop(0, 0, 2, height).resize(1, (height / 6.0).round, Magick::BoxFilter)
  end
  
  def right_edge_image
    @right_edge_image ||= image.crop(width - 3, 0, 2, height).resize(1, (height / 6.0).round, Magick::BoxFilter)
  end
  
  def difference(other_slice)
    [ distance(left_edge_image, other_slice.right_edge_image), distance(right_edge_image, other_slice.left_edge_image) ]
  end
  
  def join(right_slice)
    new_image = Magick::Image.new(width + right_slice.width, height)
    new_image.composite!(image, 0, 0, Magick::OverCompositeOp)
    new_image.composite!(right_slice.image, width, 0, Magick::OverCompositeOp)
    @image, @right_edge_image = new_image, right_slice.right_edge_image
    @width += right_slice.width
    self
  end
  
private
  
  def distance(a, b)
    max = Magick::QuantumRange.to_f
    d_sq = 0
    a.rows.times do |y|
      px_a, px_b = a.pixel_color(0, y), b.pixel_color(0, y)
      [:red, :green, :blue].each do |color|
        d_sq += (Math.log((px_a.send(color) + 1.0) / max) - Math.log((px_b.send(color) + 1.0) / max)) ** 2
      end
    end
    Math.sqrt(d_sq)
  end
  
  def initialize(img, w, h)
    @image, @width, @height = img, w, h
  end
end


# generate slices from input image
#
image = Magick::ImageList.new(ARGV[0])
w, h = image.columns, image.rows
slices = []
(w / 32.0).ceil.times do |i|
  slices << Slice.new(image.crop(i * 32, 0, 32, h, true), 32, h)
end


# unshred slices by repeatedly joining the 2 best matches until only 1 slice remains
#
loop do
  break if slices.size == 1
  
  # find the best neighboring slices for each slice
  #
  closest_neighbors = []
  slices.each_with_index do |slice, i|
    best_right_neighbor, best_left_neighbor = nil, nil
    
    slices.each_with_index do |other_slice, j|
      next if i == j
      d_right_neighbor, d_left_neighbor = slice.difference(other_slice)
      best_right_neighbor = [ d_right_neighbor, j ] if best_right_neighbor.nil? or d_right_neighbor < best_right_neighbor.first
      best_left_neighbor = [ d_left_neighbor, j ] if best_left_neighbor.nil? or d_left_neighbor < best_left_neighbor.first
    end
    
    if best_right_neighbor.first < best_left_neighbor.first
      closest_neighbors << [ best_right_neighbor.first, best_right_neighbor.last, i ]
    else
      closest_neighbors << [ best_left_neighbor.first, i, best_left_neighbor.last ]
    end
  end
  
  
  # combine the 2 best matching slices
  #
  distance, left_index, right_index = closest_neighbors.sort {|a,b| a[0] <=> b[0] }.first
  new_slices = [ slices[left_index].join(slices[right_index]) ]
  slices.each_with_index do |slice, i|
    next if i == left_index or i == right_index
    new_slices << slice
  end
  
  slices = new_slices
end

slices.first.image.write(ARGV[1])