implement iqdb wrapper

.gitignore

@@ -13,3 +13,4 @@ coverage
 *.swp
 tmp/*.jpg
 tmp/*.png
+danbooru.db

app/assets/javascripts/uploads.js

@@ -13,6 +13,14 @@
       $("#related-tags-button").trigger("click");
       $("#find-artist-button").trigger("click");
     }
+
+    if ($("#iqdb-similar").length) {
+      this.initialize_iqdb_source();
+    }
+  }
+
+  Danbooru.Upload.initialize_iqdb_source = function() {
+    $.get("/iqdb/similar_by_source", {"source": $("#upload_source").val()}).done(function(html) {$("#iqdb-similar").html(html)});
   }
 
   Danbooru.Upload.initialize_enter_on_tags = function() {

app/controllers/iqdb_controller.rb

@@ -0,0 +1,8 @@
+class IqdbController < ApplicationController
+  def similar_by_source
+    @download = Iqdb::Download.new(params[:source])
+    @download.download_from_source
+    @download.find_similar
+    render :layout => false
+  end
+end

app/logical/iqdb/download.rb

@@ -0,0 +1,21 @@
+module Iqdb
+  class Download
+    attr_reader :source, :download, :similar_posts
+
+    def initialize(source)
+      @source = source
+    end
+
+    def download_from_source
+      tempfile = Tempfile.new("iqdb-#{$PROCESS_ID}")
+      @download = Downloads::File.new(source, tempfile.path)
+      @download.download!
+    end
+
+    def find_similar
+      if Danbooru.config.iqdb_hostname_and_port
+        @similar_posts = Iqdb::Server.new(*Danbooru.config.iqdb_hostname_and_port).query(0, 3, @download.file_path)
+      end
+    end
+  end
+end

app/logical/iqdb/importer.rb

@@ -0,0 +1,13 @@
+module Iqdb
+  class Importer
+    def import!
+      IO.popen("iqdb add #{Rails.root}/danbooru.db", "w") do |io|
+        Post.find_each do |post|
+          if File.exists?(post.preview_file_path)
+            io.puts "#{post.id.to_s(16)}:#{post.preview_file_path}"
+          end
+        end
+      end
+    end
+  end
+end

app/logical/iqdb/responses/base.rb

@@ -0,0 +1,6 @@
+module Iqdb
+  module Responses
+    class Base
+    end
+  end
+end

app/logical/iqdb/responses/collection.rb

@@ -0,0 +1,27 @@
+module Iqdb
+  module Responses
+    class Collection
+      def initialize(response_string)
+        @responses = response_string.split(/\n/).map do |string|
+          ::Iqdb::Responses.const_get("Response_#{string[0..2]}").new(string[4..-1])
+        end
+      end
+
+      def matches
+        @matches ||= responses.select {|x| x.is_a?(Iqdb::Responses::Response_200) && x.score >= 0.9}
+      end
+
+      def empty?
+        matches.empty?
+      end
+
+      def errored?
+        errors.any?
+      end
+
+      def errors
+        @errors ||= responses.select {|x| x.is_a?(Iqdb::Responses::Error)}.map {|x| x.to_s}
+      end
+    end
+  end
+end

app/logical/iqdb/responses/error.rb

@@ -0,0 +1,6 @@
+module Iqdb
+  module Responses
+    class Error < Base
+    end
+  end
+end

app/logical/iqdb/responses/response_000.rb

@@ -0,0 +1,8 @@
+module Iqdb
+  module Responses
+    class Response_000 < Base
+      def initialize(response_string)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_100.rb

@@ -0,0 +1,11 @@
+module Iqdb
+  module Responses
+    class Response_100 < Base
+      attr_reader :message
+
+      def initialize(response_string)
+        @message = response_string
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_101.rb

@@ -0,0 +1,11 @@
+module Iqdb
+  module Responses
+    class Response_101 < Base
+      attr_reader :key, :value
+
+      def initialize(response_string)
+        @key, @value = response_string.split(/\=/)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_102.rb

@@ -0,0 +1,11 @@
+module Iqdb
+  module Responses
+    class Response_102 < Base
+      attr_reader :dbid, :filename
+
+      def initialize(response_string)
+        @dbid, @filename = response_string.split(/ /)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_200.rb

@@ -0,0 +1,18 @@
+module Iqdb
+  module Responses
+    class Response_200 < Base
+      attr_reader :imgid, :score, :width, :height
+
+      def initialize(response_string)
+        @imgid, @score, @width, @height = response_string.split(/ /)
+        @score = score.to_f
+        @width = width.to_i
+        @height = height.to_i
+      end
+
+      def post_id
+        imgid.to_i(16)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_201.rb

@@ -0,0 +1,19 @@
+module Iqdb
+  module Responses
+    class Response_201 < Base
+      attr_reader :dbid, :imgid, :score, :width, :height
+
+      def initialize(response_string)
+        @dbid, @imgid, @score, @width, @height = response_string.split(/ /)
+        @dbid = dbid.to_i
+        @score = score.to_f
+        @width = width.to_i
+        @height = height.to_i
+      end
+
+      def post_id
+        imgid.to_i(16)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_202.rb

@@ -0,0 +1,19 @@
+module Iqdb
+  module Responses
+    class Response_202 < Base
+      attr_reader :original_id, :stddev, :dupes
+
+      def initialize(response_string)
+        response_string =~ /^(\d+)=([0-9.]+)/
+        @original_id = $1
+        @stddev = $2
+
+        @dupes = response_string.scan(/(\d+):([0-9.]+)/).map {|x| [x[0].to_i(16), x[1].to_f]}
+      end
+
+      def original_post_id
+        original_id.to_i(16)
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_300.rb

@@ -0,0 +1,15 @@
+module Iqdb
+  module Responses
+    class Response_300 < Error
+      attr_reader :message
+
+      def initialize(response_string)
+        @message = response_string
+      end
+
+      def to_s
+        "Error: #{message}"
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_301.rb

@@ -0,0 +1,17 @@
+module Iqdb
+  module Responses
+    class Response_301 < Error
+      attr_reader :exception, :description
+
+      def initialize(response_string)
+        response_string =~ /^(\S+) (.+)/
+        @exception = $1
+        @description = $2
+      end
+
+      def to_s
+        "Exception: #{exception}: #{description}"
+      end
+    end
+  end
+end

app/logical/iqdb/responses/response_302.rb

@@ -0,0 +1,17 @@
+module Iqdb
+  module Responses
+    class Response_302 < Error
+      attr_reader :exception, :description
+
+      def initialize(response_string)
+        response_string =~ /^(\S+) (.+)/
+        @exception = $1
+        @description = $2
+      end
+
+      def to_s
+        "Fatal Exception: #{exception}: #{description}"
+      end
+    end
+  end
+end

app/logical/iqdb/server.rb

@@ -0,0 +1,56 @@
+module Iqdb
+  class Server
+    FLAG_SKETCH = 0x01
+    FLAG_GRAYSCALE = 0x02
+    FLAG_WIDTH_AS_SET = 0x08
+    FLAG_DISCARD_COMMON_COEFFS = 0x16
+
+    attr_reader :hostname, :port
+
+    def self.import(database)
+      IO.popen("iqdb #{database}", "w") do |io|
+        Post.find_each do |post|
+          puts "Adding #{post.id}"
+          io.puts "#{post.id.to_s(16)} :#{post.preview_file_path}"
+        end
+      end
+    end
+
+    def self.add(database, image_id, filename)
+      image_id_hex = image_id.to_s(16)
+      `iqdb add #{database} #{image_id_hex} :#{filename}`
+    end
+
+    def self.remove(database, image_id)
+      image_id_hex = image_id.to_s(16)
+      `iqdb remove 0 #{image_id_hex} #{database}`
+    end
+
+    def initialize(hostname, port)
+      @hostname = hostname
+      @port = port
+    end
+
+    def open
+      @socket = TCPSocket.new(hostname, port)
+    end
+
+    def close
+      @socket.close
+    end
+
+    def request
+      open
+      yield
+    ensure
+      close
+    end
+
+    def query(dbid, results, filename, flags = FLAG_DISCARD_COMMON_COEFFS)
+      request do
+        @socket.puts "query #{dbid} #{flags} #{results} #{filename}"
+        responses = Responses::Collection.new(@socket.read)
+      end
+    end
+  end
+end

app/models/post.rb

@@ -9,6 +9,8 @@ class Post < ActiveRecord::Base
   after_save :create_version
   after_save :update_parent_on_save
   after_save :apply_post_metatags, :on => :create
+  # after_save :update_iqdb, :on => :create
+  # after_destroy :remove_iqdb
   before_save :merge_old_changes
   before_save :normalize_tags
   before_save :update_tag_post_counts
@@ -1268,6 +1270,20 @@ class Post < ActiveRecord::Base
       end
     end
   end
+
+  module IqdbMethods
+    def update_iqdb
+      Danbooru.config.all_server_hosts.each do |host|
+        Iqdb::Server.delay(:queue => host).add(Danbooru.config.iqdb_file, id, preview_file_path)
+      end
+    end
+
+    def remove_iqdb
+      Danbooru.config.all_server_hosts.each do |host|
+        Iqdb::Server.delay(:queue => host).remove(Danbooru.config.iqdb_file, id)
+      end
+    end
+  end
   
   include FileMethods
   include ImageMethods

app/views/iqdb/similar_by_source.html.erb

@@ -0,0 +1,6 @@
+<% if @download.similar_posts %>
+  <h3>Similar</h3>
+  <% @download.similar_posts.each do |similar| %>
+    <%= PostPresenter.preview(Post.find(similar.post_id)) %>
+  <% end %>
+<% end %>

app/views/uploads/new.html.erb

@@ -59,6 +59,12 @@
         <%= f.text_field :parent_id %>
       </div>
 
+      <% if false && Danbooru.config.iqdb_hostname_and_port && params[:url] %>
+        <div class="input" id="iqdb-similar">
+          <p><em>Loading similar...</em></p>
+        </div>
+      <% end %>
+
       <div class="input">
         <div>
           <%= f.label :tag_string, "Tags" %>

config/danbooru_default_config.rb

@@ -300,5 +300,15 @@ module Danbooru
     def enable_dimension_autotagging
       true
     end
+
+    def iqdb_hostname_and_port
+      # ["localhost", 4000]
+      nil
+    end
+
+    def iqdb_file
+      # /var/www/danbooru2/shared/iqdb.db
+      nil
+    end
   end
 end

script/fixes/029_iqdb_import.rb

@@ -0,0 +1,5 @@
+#!/usr/bin/env ruby
+
+require File.expand_path(File.join(File.dirname(__FILE__), '..', '..', 'config', 'environment'))
+
+Iqdb::Server.import("/var/www/danbooru2/shared/iqdb.db")