iqdb: update API client to use new version of IQDB.
Replace the old IQDB API client with a new client for the new forked version of IQDB at https://github.com/danbooru/iqdb. Changes: * The /iqdb_queries endpoint now returns `hash` and `signature` fields. The `signature` is the full decoded Haar signature, while the `hash` is a encoded version of the signature. * The /iqdb_queries endpoint no longer returns `width` and `height` fields in the response (these were always 128x128). * We no longer need the IQDBs frontend server, now we talk to the IQDB instance directly. * We no longer send add/remove image commands to IQDB through AWS SQS, now we send them to IQDB directly. They are sent in a delayed job so that if IQDB is down, uploading images is still possible, the add image commands will just get queued up. * Fix a bug where regenerating an image's thumbnails didn't regenerate IQDB, because IQDB silently ignored add image commands when the image already existed in the database.
This commit is contained in:
@@ -5,7 +5,8 @@ class IqdbQueriesController < ApplicationController
|
||||
# XXX allow bare search params for backwards compatibility.
|
||||
search_params.merge!(params.slice(:url, :image_url, :file_url, :post_id, :limit, :similarity, :high_similarity).permit!)
|
||||
|
||||
@high_similarity_matches, @low_similarity_matches, @matches = IqdbProxy.new.search(search_params)
|
||||
iqdb_params = search_params.to_h.symbolize_keys
|
||||
@high_similarity_matches, @low_similarity_matches, @matches = IqdbClient.new.search(**iqdb_params)
|
||||
respond_with(@matches, template: "iqdb_queries/show")
|
||||
end
|
||||
|
||||
|
||||
@@ -31,7 +31,7 @@ class MockServicesController < ApplicationController
|
||||
render json: @data
|
||||
end
|
||||
|
||||
def iqdbs_similar
|
||||
def iqdb_query
|
||||
@data = posts.map { |post| { post_id: post.id, score: rand(0..100)} }
|
||||
render json: @data
|
||||
end
|
||||
|
||||
5
app/jobs/iqdb_add_post_job.rb
Normal file
5
app/jobs/iqdb_add_post_job.rb
Normal file
@@ -0,0 +1,5 @@
|
||||
class IqdbAddPostJob < ApplicationJob
|
||||
def perform(post)
|
||||
IqdbClient.new.add_post(post)
|
||||
end
|
||||
end
|
||||
5
app/jobs/iqdb_remove_post_job.rb
Normal file
5
app/jobs/iqdb_remove_post_job.rb
Normal file
@@ -0,0 +1,5 @@
|
||||
class IqdbRemovePostJob < ApplicationJob
|
||||
def perform(post_id)
|
||||
IqdbClient.new.remove(post_id)
|
||||
end
|
||||
end
|
||||
86
app/logical/iqdb_client.rb
Normal file
86
app/logical/iqdb_client.rb
Normal file
@@ -0,0 +1,86 @@
|
||||
class IqdbClient
|
||||
class Error < StandardError; end
|
||||
attr_reader :iqdb_url, :http
|
||||
|
||||
def initialize(iqdb_url: Danbooru.config.iqdb_url.to_s, http: Danbooru::Http.new)
|
||||
@iqdb_url = iqdb_url.chomp("/")
|
||||
@http = http
|
||||
end
|
||||
|
||||
concerning :QueryMethods do
|
||||
def search(post_id: nil, file: nil, url: nil, image_url: nil, file_url: nil, similarity: 0.0, high_similarity: 65.0, limit: 20)
|
||||
limit = limit.to_i.clamp(1, 1000)
|
||||
similarity = similarity.to_f.clamp(0.0, 100.0)
|
||||
high_similarity = high_similarity.to_f.clamp(0.0, 100.0)
|
||||
|
||||
if file.blank?
|
||||
if url.present?
|
||||
file = download(url, :preview_url)
|
||||
elsif image_url.present?
|
||||
file = download(image_url, :url)
|
||||
elsif file_url.present?
|
||||
file = download(file_url, :image_url)
|
||||
elsif post_id.present?
|
||||
file = Post.find(post_id).file(:preview)
|
||||
else
|
||||
return [[], [], []]
|
||||
end
|
||||
end
|
||||
|
||||
results = query(file, limit: limit)
|
||||
results = results.select { |result| result["score"] >= similarity }.take(limit)
|
||||
matches = decorate_posts(results)
|
||||
high_similarity_matches, low_similarity_matches = matches.partition { |match| match["score"] >= high_similarity }
|
||||
|
||||
[high_similarity_matches, low_similarity_matches, matches]
|
||||
ensure
|
||||
file.try(:close)
|
||||
end
|
||||
|
||||
def download(url, type)
|
||||
strategy = Sources::Strategies.find(url)
|
||||
download_url = strategy.send(type)
|
||||
file = strategy.download_file!(download_url)
|
||||
file
|
||||
end
|
||||
|
||||
def decorate_posts(json)
|
||||
post_ids = json.map { |match| match["post_id"] }
|
||||
posts = Post.where(id: post_ids).group_by(&:id).transform_values(&:first)
|
||||
|
||||
json.map do |match|
|
||||
post = posts.fetch(match["post_id"], nil)
|
||||
match.with_indifferent_access.merge(post: post) if post
|
||||
end.compact
|
||||
end
|
||||
end
|
||||
|
||||
def add_post(post)
|
||||
return unless post.has_preview?
|
||||
preview_file = post.file(:preview)
|
||||
add(post.id, preview_file)
|
||||
end
|
||||
|
||||
concerning :HttpMethods do
|
||||
def query(file, limit: 20)
|
||||
file = HTTP::FormData::File.new(file)
|
||||
request(:post, "query", form: { file: file }, params: { limit: limit })
|
||||
end
|
||||
|
||||
def add(post_id, file)
|
||||
file = HTTP::FormData::File.new(file)
|
||||
request(:post, "images/#{post_id}", form: { file: file })
|
||||
end
|
||||
|
||||
def remove(post_id)
|
||||
request(:delete, "images/#{post_id}")
|
||||
end
|
||||
|
||||
def request(method, url, **options)
|
||||
return [] if iqdb_url.blank? # do nothing if iqdb isn't configured
|
||||
response = http.timeout(30).send(method, "#{iqdb_url}/#{url}", **options)
|
||||
raise Error, "IQDB error: #{response.parse}" if response.status != 200
|
||||
response.parse
|
||||
end
|
||||
end
|
||||
end
|
||||
@@ -1,77 +0,0 @@
|
||||
class IqdbProxy
|
||||
class Error < StandardError; end
|
||||
attr_reader :http, :iqdbs_server
|
||||
|
||||
def initialize(http: Danbooru::Http.new, iqdbs_server: Danbooru.config.iqdbs_server)
|
||||
@iqdbs_server = iqdbs_server
|
||||
@http = http
|
||||
end
|
||||
|
||||
def enabled?
|
||||
iqdbs_server.present?
|
||||
end
|
||||
|
||||
def download(url, type)
|
||||
strategy = Sources::Strategies.find(url)
|
||||
download_url = strategy.send(type)
|
||||
file = strategy.download_file!(download_url)
|
||||
file
|
||||
end
|
||||
|
||||
def search(params)
|
||||
limit = params[:limit]&.to_i&.clamp(1, 1000) || 20
|
||||
similarity = params[:similarity]&.to_f&.clamp(0.0, 100.0) || 0.0
|
||||
high_similarity = params[:high_similarity]&.to_f&.clamp(0.0, 100.0) || 65.0
|
||||
|
||||
if params[:file].present?
|
||||
file = params[:file]
|
||||
results = query(file: file, limit: limit)
|
||||
elsif params[:url].present?
|
||||
file = download(params[:url], :preview_url)
|
||||
results = query(file: file, limit: limit)
|
||||
elsif params[:image_url].present?
|
||||
file = download(params[:image_url], :url)
|
||||
results = query(file: file, limit: limit)
|
||||
elsif params[:file_url].present?
|
||||
file = download(params[:file_url], :image_url)
|
||||
results = query(file: file, limit: limit)
|
||||
elsif params[:post_id].present?
|
||||
url = Post.find(params[:post_id]).preview_file_url
|
||||
results = query(url: url, limit: limit)
|
||||
else
|
||||
results = []
|
||||
end
|
||||
|
||||
results = results.select { |result| result["score"] >= similarity }.take(limit)
|
||||
matches = decorate_posts(results)
|
||||
high_similarity_matches, low_similarity_matches = matches.partition { |match| match["score"] >= high_similarity }
|
||||
|
||||
[high_similarity_matches, low_similarity_matches, matches]
|
||||
ensure
|
||||
file.try(:close)
|
||||
end
|
||||
|
||||
def query(file: nil, url: nil, limit: 20)
|
||||
raise NotImplementedError, "the IQDBs service isn't configured" unless enabled?
|
||||
|
||||
file = HTTP::FormData::File.new(file) if file
|
||||
form = { file: file, url: url, limit: limit }.compact
|
||||
response = http.timeout(30).post("#{iqdbs_server}/similar", form: form)
|
||||
|
||||
raise Error, "IQDB error: #{response.status}" if response.status != 200
|
||||
raise Error, "IQDB error: #{response.parse["error"]}" if response.parse.is_a?(Hash)
|
||||
raise Error, "IQDB error: #{response.parse.first}" if response.parse.try(:first).is_a?(String)
|
||||
|
||||
response.parse
|
||||
end
|
||||
|
||||
def decorate_posts(json)
|
||||
post_ids = json.map { |match| match["post_id"] }
|
||||
posts = Post.where(id: post_ids).group_by(&:id).transform_values(&:first)
|
||||
|
||||
json.map do |match|
|
||||
post = posts.fetch(match["post_id"], nil)
|
||||
match.with_indifferent_access.merge(post: post) if post
|
||||
end.compact
|
||||
end
|
||||
end
|
||||
@@ -85,6 +85,8 @@ class UploadService
|
||||
)
|
||||
end
|
||||
|
||||
@post.update_iqdb
|
||||
|
||||
upload.update(status: "completed", post_id: @post.id)
|
||||
|
||||
@post
|
||||
|
||||
@@ -128,7 +128,7 @@ class UploadService
|
||||
replacement.save!
|
||||
post.save!
|
||||
|
||||
post.update_iqdb_async
|
||||
post.update_iqdb
|
||||
end
|
||||
|
||||
def purge_cached_urls(post)
|
||||
|
||||
@@ -36,8 +36,6 @@ class Post < ApplicationRecord
|
||||
after_save :update_parent_on_save
|
||||
after_save :apply_post_metatags
|
||||
after_commit :delete_files, :on => :destroy
|
||||
after_commit :remove_iqdb_async, :on => :destroy
|
||||
after_commit :update_iqdb_async, :on => :create
|
||||
|
||||
belongs_to :approver, class_name: "User", optional: true
|
||||
belongs_to :uploader, :class_name => "User", :counter_cache => "post_upload_count"
|
||||
@@ -977,6 +975,8 @@ class Post < ApplicationRecord
|
||||
update_parent_on_destroy
|
||||
end
|
||||
end
|
||||
|
||||
remove_iqdb # this is non-transactional
|
||||
end
|
||||
|
||||
def ban!
|
||||
@@ -1319,16 +1319,15 @@ class Post < ApplicationRecord
|
||||
|
||||
def regenerate!(category, user)
|
||||
if category == "iqdb"
|
||||
update_iqdb_async
|
||||
update_iqdb
|
||||
|
||||
ModAction.log("<@#{user.name}> regenerated IQDB for post ##{id}", :post_regenerate_iqdb, user)
|
||||
else
|
||||
media_file = MediaFile.open(file, frame_data: pixiv_ugoira_frame_data)
|
||||
UploadService::Utils.process_resizes(self, nil, id, media_file: media_file)
|
||||
|
||||
# XXX This may be racy; the thumbnail may not be purged from Cloudflare by the time IQDB tries to download it.
|
||||
purge_cached_urls!
|
||||
update_iqdb_async
|
||||
update_iqdb
|
||||
|
||||
ModAction.log("<@#{user.name}> regenerated image samples for post ##{id}", :post_regenerate, user)
|
||||
end
|
||||
@@ -1340,33 +1339,15 @@ class Post < ApplicationRecord
|
||||
end
|
||||
end
|
||||
|
||||
module IqdbMethods
|
||||
extend ActiveSupport::Concern
|
||||
|
||||
module ClassMethods
|
||||
def iqdb_sqs_service
|
||||
SqsService.new(Danbooru.config.aws_sqs_iqdb_url)
|
||||
end
|
||||
|
||||
def iqdb_enabled?
|
||||
Danbooru.config.aws_sqs_iqdb_url.present?
|
||||
end
|
||||
|
||||
def remove_iqdb(post_id)
|
||||
if iqdb_enabled?
|
||||
iqdb_sqs_service.send_message("remove\n#{post_id}")
|
||||
end
|
||||
end
|
||||
concerning :IqdbMethods do
|
||||
def update_iqdb
|
||||
# performs IqdbClient.new.add_post(post)
|
||||
IqdbAddPostJob.perform_later(self)
|
||||
end
|
||||
|
||||
def update_iqdb_async
|
||||
if Post.iqdb_enabled? && has_preview?
|
||||
Post.iqdb_sqs_service.send_message("update\n#{id}\n#{preview_file_url}")
|
||||
end
|
||||
end
|
||||
|
||||
def remove_iqdb_async
|
||||
Post.remove_iqdb(id)
|
||||
def remove_iqdb
|
||||
# performs IqdbClient.new.remove(id)
|
||||
IqdbRemovePostJob.perform_later(id)
|
||||
end
|
||||
end
|
||||
|
||||
@@ -1460,7 +1441,6 @@ class Post < ApplicationRecord
|
||||
include ApiMethods
|
||||
extend SearchMethods
|
||||
include PixivMethods
|
||||
include IqdbMethods
|
||||
include ValidationMethods
|
||||
|
||||
has_bit_flags ["has_embedded_notes", "has_cropped"]
|
||||
|
||||
@@ -16,7 +16,7 @@
|
||||
</div>
|
||||
<% end %>
|
||||
|
||||
<% if Danbooru.config.iqdbs_server %>
|
||||
<% if Danbooru.config.iqdb_url %>
|
||||
<% if params[:url] %>
|
||||
<div class="input" id="iqdb-similar">
|
||||
<p><em>Loading similar...</em></p>
|
||||
|
||||
Reference in New Issue
Block a user