Add IQDB for yiffy2

Gemfile

@@ -101,3 +101,5 @@ gem "aws-sdk-s3", "~> 1.149"
 gem "opensearch-ruby", "~> 3.3"
 
 gem "good_job", "~> 3.99"
+
+gem "faraday", "~> 2.9"

Gemfile.lock

@@ -357,6 +357,7 @@ DEPENDENCIES
   debug
   dotenv-rails
   ed25519 (~> 1.3)
+  faraday (~> 2.9)
   filesize (~> 0.2.0)
   github_webhook (~> 1.4)
   good_job (~> 3.99)

app/controllers/yiff_rest/api_v2/images_controller.rb

@@ -57,6 +57,24 @@ module YiffRest
         @image = APIImage.find(params[:id])
       end
 
+      def iqdb
+      end
+
+      def query_iqdb
+        @sp = search_iqdb_params
+        @sc = (@sp[:score_cutoff].presence || 60).to_i
+        if @sp[:file].present?
+          @results = Iqdb.query_file(@sp[:file], @sc)
+        elsif @sp[:url].present?
+          @results = Iqdb.query_url(@sp[:url], @sc)
+        elsif @sp[:image_id].present?
+          @results = Iqdb.query_image(APIImage.find(@sp[:image_id]), @sc)
+        end
+        @results ||= []
+        @pagy, @images = pagy_array(@results.pluck("image"), size: [1, 2, 2, 1])
+        render(:index)
+      end
+
       private
 
       def load_category
@@ -72,6 +90,10 @@ module YiffRest
         permit_search_params(%i[category md5 original_url artist])
       end
 
+      def search_iqdb_params
+        permit_search_params(%i[file url image_id score_cutoff])
+      end
+
       def create_params
         params.fetch(:api_image, {}).permit(:category, :original_url, :sources_string, :artists_string, :file)
       end

app/jobs/iqdb_remove_job.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class IqdbRemoveJob < ApplicationJob
+  queue_as :iqdb
+
+  def perform(id)
+    Iqdb.remove_image(id)
+  end
+end

app/jobs/iqdb_update_job.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+class IqdbUpdateJob < ApplicationJob
+  queue_as :iqdb
+
+  def perform(id)
+    image = APIImage.find(id)
+    return if image.blank?
+    Iqdb.update_image(image)
+  end
+end

app/logical/iqdb.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+module Iqdb
+  class Error < StandardError; end
+
+  module_function
+
+  IQDB_NUM_PIXELS = 128
+
+  def endpoint
+    Websites.config.iqdb_server
+  end
+
+  def enabled?
+    endpoint.present?
+  end
+
+  def make_request(path, request_type, body = nil)
+    conn = Faraday.new(Websites.config.faraday_options)
+    conn.send(request_type, endpoint + path, body&.to_json, { content_type: "application/json" })
+  rescue Faraday::Error
+    raise(Error, "This service is temporarily unavailable. Please try again later.")
+  end
+
+  def update_image(image)
+    Websites.config.yiffy2_storage.get(image.path) do |file|
+      Tempfile.open("yiffy2-#{image.md5}") do |tempfile|
+        file.binmode
+        tempfile.binmode
+        tempfile.write(file.read)
+        tempfile.rewind
+        thumb = generate_thumbnail(tempfile.path)
+        raise(Error, "failed to generate thumb for #{image.iqdb_id}") unless thumb
+        response = make_request("/images/#{image.iqdb_id}", :post, get_channels_data(thumb))
+        raise(Error, "iqdb request failed") if response.status != 200
+      end
+    end
+  end
+
+  def remove_image(id)
+    response = make_request("/images/#{id}", :delete)
+    raise(Error, "iqdb request failed") if response.status != 200
+  end
+
+  def query_url(url, score_cutoff)
+    file = FileDownload.new(url).download!
+    query_file(file, score_cutoff)
+  end
+
+  def query_image(image, score_cutoff)
+    Websites.config.yiffy2_storage.get(image.path) do |file|
+      query_file(file, score_cutoff)
+    end
+  end
+
+  def query_file(file, score_cutoff)
+    thumb = generate_thumbnail(file.path)
+    return [] unless thumb
+
+    response = make_request("/query", :post, get_channels_data(thumb))
+    return [] if response.status != 200
+
+    process_iqdb_result(JSON.parse(response.body), score_cutoff)
+  end
+
+  def query_hash(hash, score_cutoff)
+    response = make_request("/query", :post, { hash: hash })
+    return [] if response.status != 200
+
+    process_iqdb_result(JSON.parse(response.body), score_cutoff)
+  end
+
+  def process_iqdb_result(json, score_cutoff)
+    raise(Error, "Server returned an error. Most likely the url is not found.") unless json.is_a?(Array)
+
+    json.filter! { |entry| (entry["score"] || 0) >= (score_cutoff.presence || 60).to_i }
+    json.map do |x|
+      x["image"] = APIImage.find_by!(iqdb_id: x["post_id"])
+      x
+    rescue ActiveRecord::RecordNotFound
+      nil
+    end.compact
+  end
+
+  def generate_thumbnail(file_path)
+    Vips::Image.thumbnail(file_path, IQDB_NUM_PIXELS, height: IQDB_NUM_PIXELS, size: :force)
+  rescue Vips::Error => e
+    Rails.logger.error("failed to generate thumbnail for #{file_path}")
+    Rails.logger.error(e)
+    nil
+  end
+
+  def get_channels_data(thumbnail)
+    r = []
+    g = []
+    b = []
+    is_grayscale = thumbnail.bands == 1
+    thumbnail.to_a.each do |data|
+      data.each do |rgb|
+        r << rgb[0]
+        g << (is_grayscale ? rgb[0] : rgb[1])
+        b << (is_grayscale ? rgb[0] : rgb[2])
+      end
+    end
+    { channels: { r: r, g: g, b: b } }
+  end
+end

app/logical/requests/bunny.rb

@@ -23,9 +23,11 @@ module Requests
       !get(path).nil?
     end
 
-    def get(path)
+    def get(path, &block)
       r = self.class.get("/#{storage_zone_name}/#{path}", headers: { "AccessKey" => access_key })
-      r.success? ? r.body : nil
+      file = r.success? ? r.body : nil
+      block&.call(file)
+      file
     end
 
     def put(path, body)

app/logical/storage_manager/local.rb

@@ -19,8 +19,8 @@ module StorageManager
       File.exist?("#{base_path}#{path}")
     end
 
-    def get(path)
-      File.open("#{base_path}#{path}", "r", binmode: true)
+    def get(path, &)
+      File.open("#{base_path}#{path}", "r", binmode: true, &)
     end
 
     def put(path, io)

app/logical/storage_manager/s3.rb

@@ -22,8 +22,10 @@ module StorageManager
       @s3.object(trim(path)).exists?
     end
 
-    def get(path)
-      @s3.object(trim(path)).get.body
+    def get(path, &block)
+      file = @s3.object(trim(path)).get.body
+      block&.call(file)
+      file
     end
 
     def put(path, io)

app/models/api_image.rb

@@ -18,9 +18,11 @@ class APIImage < ApplicationRecord
   after_create :invalidate_cache
   after_create :send_created, unless: :no_webhook_messages
   after_create :initialize_short_url
+  after_create :update_iqdb
   after_update :update_files, if: :saved_change_to_category?
   after_update :send_updated, unless: :no_webhook_messages
   after_destroy :delete_files
+  after_destroy :remove_iqdb
   after_destroy :send_deleted, unless: :no_webhook_messages
 
   def delete_files
@@ -323,4 +325,32 @@ class APIImage < ApplicationRecord
   end
 
   include WebhookMethods
+
+  def is_image?
+    %w[png jpg].include?(file_ext)
+  end
+
+  def is_gif?
+    file_ext == "gif"
+  end
+
+  def update_iqdb
+    return unless is_image?
+    IqdbUpdateJob.perform_later(id)
+  end
+
+  def update_iqdb!
+    return unless is_image?
+    IqdbUpdateJob.perform_now(id)
+  end
+
+  def remove_iqdb
+    return unless is_image?
+    IqdbRemoveJob.perform_later(iqdb_id)
+  end
+
+  def remove_iqdb!
+    return unless is_image?
+    IqdbRemoveJob.perform_now(iqdb_id)
+  end
 end

app/views/yiff_rest/api_v2/images/index.html.erb

@@ -51,6 +51,7 @@
           <%= link_to "Edit", edit_yiff_rest_api_v2_manage_image_path(manage_id: params[:id], id: img.id) %> |
           <%= link_to "Delete", yiff_rest_api_v2_manage_image_path(manage_id: params[:id], id: img.id), method: :delete %><br>
           <%= link_to "Delete W/R", delete_with_reason_yiff_rest_api_v2_manage_image_path(manage_id: params[:id], id: img.id) %>
+          <%= link_to("Find Similar", iqdb_query_yiff_rest_api_v2_manage_images_path(search: { image_id: img.id })) %>
         </td>
       </tr>
     <% end %>

app/views/yiff_rest/api_v2/images/iqdb.html.erb

@@ -0,0 +1,16 @@
+<% content_for(:page_title) do %>
+  YiffyAPI - IQDB
+<% end %>
+
+<%= render "yiff_rest/api_v2/manage/nav" %>
+
+<div class="w-100">
+  <div style="width: 40%; margin-left: 30%;">
+    <%= simple_form_for(:search, url: iqdb_query_yiff_rest_api_v2_manage_images_path, method: :post) do |f| %>
+      <%= f.input(:file, as: :file) %>
+      <%= f.input(:url, label: "File URL") %>
+      <%= f.input(:image_id, label: "Image ID") %>
+      <%= f.button(:submit, "Query", name: nil) %>
+    <% end %>
+  </div>
+</div>

app/views/yiff_rest/api_v2/manage/_nav.html.erb

@@ -27,6 +27,9 @@
               </li>
             <% end %>
           <% end %>
+          <li class="nav-item">
+            <%= link_to("IQDB", iqdb_query_yiff_rest_api_v2_manage_images_path, class: "nav-link") %>
+          </li>
         <% end %>
       </ul>
       <div class="d-flex">

config/default_config.rb

@@ -56,6 +56,10 @@ module Websites
       "opensearch.websites4.containers.local"
     end
 
+    def iqdb_server
+      "http://iqdb.websites4.containers.local:5588"
+    end
+
     def github_webhook_secret
     end
 
@@ -231,6 +235,17 @@ module Websites
       }
     end
 
+    # https://lostisland.github.io/faraday/#/customization/connection-options
+    def faraday_options
+      {
+        request: {
+          timeout:      10,
+          open_timeout: 5,
+        },
+        headers: http_headers,
+      }
+    end
+
     def yiffy2_cdn_url
       return "http://yiffy2.local" if Rails.env.development?
       "https://v2.yiff.media"

config/initializers/pagy.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
 
 require "pagy/extras/bootstrap"
+require "pagy/extras/array"
 Pagy::DEFAULT[:items] = 20

config/routes/yiff_rest_routes.rb

@@ -54,6 +54,11 @@ module YiffRestRoutes
                 put(:sync)
                 resources(:images, as: "manage_images") do
                   get(:delete_with_reason, on: :member)
+                  collection do
+                    get(:iqdb)
+                    get("/iqdb/query", action: :query_iqdb)
+                    post("/iqdb/query", action: :query_iqdb)
+                  end
                 end
               end
             end

db/fixes/001_iqdb_import.rb

@@ -0,0 +1,6 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require File.expand_path(File.join(File.dirname(__FILE__), "..", "..", "config", "environment"))
+
+APIImage.find_each(&:update_iqdb)

db/migrate/20241013040658_add_iqdb_id_to_api_images.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+class AddIqdbIdToAPIImages < ActiveRecord::Migration[7.1]
+  def change
+    add_column(:api_images, :iqdb_id, :bigserial, null: false) # rubocop:disable Rails/NotNullColumn
+  end
+end

db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[7.1].define(version: 2024_08_06_104838) do
+ActiveRecord::Schema[7.1].define(version: 2024_10_13_040658) do
   create_schema "e621"
 
   # These are extensions that must be enabled in order to support this database
@@ -58,6 +58,7 @@ ActiveRecord::Schema[7.1].define(version: 2024_08_06_104838) do
     t.integer "file_size", null: false
     t.datetime "created_at", null: false
     t.datetime "updated_at", null: false
+    t.bigserial "iqdb_id", null: false
     t.index ["creator_id"], name: "index_api_images_on_creator_id"
   end
 

docker-compose.yml

@@ -123,8 +123,20 @@ services:
     hostname: rethinkdb.websites4.containers.local
     labels:
       - "hostname=rethinkdb.websites4.containers.local"
-    networks:
-      - default
+
+  iqdb:
+    image: ghcr.io/e621ng/iqdb:ad5e363879d4e4b80a0821253ae610b7859c5d32
+    command: iqdb http 0.0.0.0 5588 /iqdb/iqdb.db
+    healthcheck:
+      test: wget --no-verbose --tries=1 --spider http://127.0.0.1:5588/status || exit 1
+      interval: 30s
+      timeout: 2s
+      retries: 5
+    volumes:
+      - iqdb_data:/iqdb
+    hostname: iqdb.websites4.containers.local
+    labels:
+      - "hostname=iqdb.websites4.containers.local"
 
 networks:
   default:
@@ -145,3 +157,4 @@ volumes:
   yiffy2_image_data:
   oceanic_docs_data:
   rethinkdb_data:
+  iqdb_data:

lib/middleware/dev_host.rb

@@ -2,7 +2,7 @@
 
 module Middleware
   class DevHost
-    DEFAULT_HOST = "admin.furry.computer"
+    DEFAULT_HOST = "v2.yiff.rest"
     def initialize(app)
       @app = app
     end
@@ -10,6 +10,7 @@ module Middleware
     def call(env)
       request = Rack::Request.new(env)
       domain = request.params["domain"].presence || DEFAULT_HOST
+      domain = OtherRoutes::ADMIN_DOMAIN if %w[/jobs /exceptions].any? { |p| request.path.start_with?(p) }
 
       if Rails.env.development? && domain
         env["websites.dev_host"] = domain