Add IQDB for yiffy2

This commit is contained in:
Donovan Daniels 2024-10-13 01:04:55 -05:00
parent da068e44d0
commit 624ac2f738
Signed by: Donovan_DMC
GPG Key ID: 907D29CBFD6157BA
21 changed files with 265 additions and 10 deletions

View File

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

View File

@ -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)

View File

@ -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

View File

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

View File

@ -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

107
app/logical/iqdb.rb Normal file
View File

@ -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

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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

View File

@ -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 %>

View File

@ -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>

View File

@ -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">

View File

@ -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"

View File

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

View File

@ -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

6
db/fixes/001_iqdb_import.rb Executable file
View File

@ -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)

View File

@ -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

3
db/schema.rb generated
View File

@ -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

View File

@ -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:

View File

@ -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