Instagram Engineering Challengeを解いてみる

http://instagram-engineering.tumblr.com/post/12651721845/instagram-engineering-challenge-the-unshredder

以上のようにバラバラに切断された画象を以下のように戻せ、という問題。

画象の分割数は予め与えられれているものとする。(自動的に画象の分割数を求めるとボーナス得点有り)


とりあえず隣接するカラムの類似度(画素間の距離)を行単位で求め、類似度が高いものから順に並べていけばいいんじゃないかと考えた。
画素間の距離をRGBで求めるとうまく行かなかったけど、HSVで求めるとうまく行った。ふむ。

# -- coding utf-8 --
from PIL import Image, ImageChops, ImageStat
import math, operator, sys, colorsys

def calc_diff(p1, p2):
  p1 = colorsys.rgb_to_hsv(p1[0] / 255., p1[1] / 255., p1[2] / 255.)
  p2 = colorsys.rgb_to_hsv(p2[0] / 255., p2[1] / 255., p2[2] / 255.)
  
  return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2 + (p1[2] - p2[2]) ** 2)

def find_match_scores(image, width, height, columns, shred_width):
  # 行単位での画素の距離の総和をもとめる
  # 値が大きいほど類似度が低い
  scores = {}
  for i in range(columns):
    right1 = shred_width * (i + 1) - 1
    left1 = shred_width * i
    next_left = None
    score = None  

    for j in range(columns):
      if i == j:
        continue
        
      right2 = shred_width * (j + 1) - 1
      left2 = shred_width * j

      tmp_score = 0.0
      for h in range(height):
        p1 = image.getpixel((right1, h))
        p2 = image.getpixel((left2, h))
        tmp_score += calc_diff(p1, p2)
      
      if score is None or score > tmp_score:
        score = tmp_score
        next_left = left2
     
    scores[left1] = (next_left, score)
  
  return scores

def find_rightest_colum(image, width, height, scores):
  right = None
  # 遷移先が重複しているものがある場合、スコアが高いほうが右はじのカラムだと仮定する
  for k, v in scores.items():
    for k2, v2 in scores.items():
      if k == k2:
        continue
      if v[0] == v2[0]:
        if v[1] > v2[1]:
          right = k
        else:
          right = k2
        return right
  # 上の仮定に該当するものがない場合はスコアが一番たかいものを右はじのカラムだと仮定する
  if right is None:
    max_item = sorted(scores.items(), key = lambda x : x[1])[0]
    return  max_item[0]

def find_order(rightest_column, scores):
  order = []
  # 右はじのカラムにたどりくまでを逆算してもとめるため、key-valueの関係を逆転させる
  reverse_scores = {}
  for k, v in scores.items():
    reverse_scores[v[0]] = k
  
  current = rightest_column
  while True:
    order.append(current)
    if current not in reverse_scores:
      break
    current = reverse_scores[current]
  # 左から右の順に変更
  order.reverse()  
  return order

def paint_out(image, width, height, shred_width, order):
  output = Image.new('RGBA', (width, height))
  for i in range(len(order)):
    source_box = (order[i], 0, order[i] + shred_width, height)
    column = image.crop(source_box)
    destination_box = (i * shred_width, 0, (i + 1) * shred_width, height)
    output.paste(column, destination_box)
  return output
 
if __name__ == '__main__':
  image = Image.open("TokyoPanoramaShredded.png")
  width, height = image.size
  columns = 20
  shred_width = shred_width = width / columns
  # 各カラム間の類似度を計算する
  scores = find_match_scores(image, width, height, columns, shred_width)
  # 右はじのカラムを求める
  rightest_column = find_rightest_colum(image, width, height, scores) 
  del(scores[rightest_column])
  # 並び順を求める
  order = find_order(rightest_column, scores)
  
  output_image = paint_out(image, width, height, shred_width, order)
  output_image.save("TokyoPanoramaUnshredded.png")