sift test

vlfeat.rb

#!/usr/bin/env ruby

require 'rubygems'
require 'rmagick'

include Math

class Node  
  attr_accessor :location, :left_child, :right_child

  def initialize
    @location = nil
    @left_child = nil
    @right_child = nil
  end

  def is_leaf
    return (self.left_child.nil? and self.right_child.nil?)
  end
end

def square_distance(ptA, ptB)
  dist = 0
  dims = Range.new(0, ptA.size, true)
  for dim in dims do
    dist += (ptA[dim] - ptB[dim])**2
  end
  return dist
end

def closer(ref, ptA, ptB)
  distA = square_distance(ref, ptA)
  distB = square_distance(ref, ptB)
  return distA < distB ? ptA : ptB
end

def process_image(imagename, resultname, params="--edge-thresh 10 --peak-thresh 5")
  img = Magick::Image.from_blob( File.read(imagename) ).shift
  if img.format() != "PGM" 
    img.format = "PGM"
    File.open("tmp.pgm", "wb") {|f| f << img.to_blob}
    imagename = "tmp.pgm"
  end

  cmmd = "sift " + imagename + " --output=" + resultname + " " + params
  system(cmmd)
  puts "processed " + imagename + " to " + resultname
end

def read_features_from_file(filename)
  keypt = Array.new
  feat = Array.new
  IO.foreach(filename) do |s|
    keypt.push( s.split(" ")[0..3] )
    feat.push( s.split(" ")[4..-1] )
  end
  return keypt, feat
end

def draw_circle(im, pt)
  gc = Magick::Draw.new
  gc.fill_opacity(0)
  gc.stroke('red')

  origin_x = pt[0].to_i
  origin_y = pt[1].to_i
  scl = pt[2].to_f
  perim_x = (pt[0].to_f - scl * cos(0.0)).to_i
  perim_y = origin_y
  gc.circle(origin_x, origin_y, perim_x, perim_y)
  gc.draw(im)
end

def draw_orientation(im, pt)
  scale = 5.0
  gc = Magick::Draw.new
  gc.stroke('blue')

  here_x = pt[0].to_i
  here_y = pt[1].to_i
  scl = pt[2].to_f
  ori = pt[3].to_f
  len = scl * cos(0.0)  
  there_x = here_x + (len * cos(ori)).to_i
  there_y = here_y + (len * sin(ori)).to_i
  gc.line(here_x, here_y, there_x, there_y)
  gc.draw(im)
end

def plot_features(im, keypt)

  for pt in keypt
    # circle
    draw_circle(im, pt)

    # orientation
    draw_orientation(im, pt)
  end
  
end

def kdtree_build(pointList, depth=0)
  if pointList.size == 0 then
    return nil
  end

  dim = pointList[0].size
  axis = depth % dim
  sortedList = pointList.sort{|p,q| p[axis] <=> q[axis]}
  median = sortedList.size / 2
  
  node = Node.new
  node.location = sortedList[median]
  node.left_child = kdtree_build(sortedList[0, median], depth+1)
  node.right_child = kdtree_build(sortedList[median+1..-1], depth+1)  
end

def kdtree_search(qt, node, depth=0)
  if node.nil? then
    return nil
  end

  # if we have reached a leaf
  if node.is_leaf then
    return node.location
  end

  # this node is no leaf

  dim = node.location.size
  axis = depth % dim

  # compare query point and point of current node in selected dimension
  if qt[axis] < node.location[axis] then

    if node.left_child then
      if node.left_child.is_leaf then
        leaf = node.left_child.location
      else
        leaf = kdtree_search(qt, node.left_child, depth+1)
      end
    else
      leaf = kdtree_search(qt, node.right_child, depth+1)
    end

    if node.left_child && node.right_child && 
        square_distance(qt, leaf) > (qt[axis]-node.location[axis])**2 then
      leaf = closer(qt, leaf, kdtree_search(qt, node.right_child, depth+1))
    end
                    
  else
    if node.right_child then
      if node.right_child.is_leaf then
        leaf = node.right_child.location
      else
        leaf = kdtree_search(qt, node.right_child, depth+1)
      end
    else
      leaf = kdtree_search(qt, node.left_child, depth+1)
    end

    if node.right_child && node.left_child &&
        square_distance(qt, leaf) > (qt[axis]-node.location[axis])**2 then
      leaf = closer(qt, leaf, kdtree_search(qt, node.left_child, depth+1))
    end
    
  end

  return closer(qt, leaf, node.location)
end

def match_kdtree(feat, feat2)
  # build a kd tree
  kdtree = kdtree_build(feat2)

  # Finds an image feature's nearest neighbors in a kd tree
  m = Array.new
  for qt in feat do
    mpt = kdtree_search(qt, kdtree)
    m.push(mpt)
  end

  return m
end

if __FILE__ == $0
  process_image('box.pgm', 'box.sift')
  keypt, feat = read_features_from_file('box.sift')

  im = Magick::Image.read('box.pgm').first
  plot_features(im, keypt)
  im.display

  process_image('scene.pgm', 'scene.sift')
  keypt2, feat2 = read_features_from_file('scene.sift')
  im2 = Magick::Image.read('scene.pgm').first

  m = match_kdtree(feat, feat2)
  print "matched: ", m.size, " pts\n"

  
end