From 1bb383703d07a5043b4ca02cc3efe57b4cae4c3e Mon Sep 17 00:00:00 2001 From: evazion Date: Thu, 2 Sep 2021 03:34:50 -0500 Subject: [PATCH] iqdb: allow searching images by iqdb hash. Example: https://danbooru.donmai.us/iqdb_queries?search[hash]=iqdb_3fe4c0bd5a88fab53fbe5104efdc43ba3fa094416e6e039cf880f9f1fafafafffc80fd7bfd7dfd80fe72fe79fe80fefafefdfeffff75ff79ff7aff7dfff1fffb000100020003000400080082008a00940184018802000287028c028f030004000600078308001080ef80f800f8fdf900fafdfc00fc7dfcfafe00fe72fefdff7effdffff0fff8fffcfffd0007000f001f003e0080008f0106018d02800282038003810403040806800683070b078007870d800d830f801f00f800f96cfa76fb00fb7efbf4fc00fcfdfd00fd74fdfbfe78fe7efef5fef9fefbfefeff6cff76ff7efffcfffe00070080008300860087008b01030106018002820283028503800402040a0502058b0d80 The hash may be obtained from a regular IQDB search, or by calculating it yourself (an exercise for the reader). --- app/controllers/iqdb_queries_controller.rb | 2 +- app/logical/iqdb_client.rb | 41 ++++++++++++++-------- app/views/iqdb_queries/show.html.erb | 1 + 3 files changed, 29 insertions(+), 15 deletions(-) diff --git a/app/controllers/iqdb_queries_controller.rb b/app/controllers/iqdb_queries_controller.rb index 408f390f4..f7d529256 100644 --- a/app/controllers/iqdb_queries_controller.rb +++ b/app/controllers/iqdb_queries_controller.rb @@ -3,7 +3,7 @@ class IqdbQueriesController < ApplicationController def show # 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!) + search_params.merge!(params.slice(:url, :hash, :image_url, :file_url, :post_id, :limit, :similarity, :high_similarity).permit!) iqdb_params = search_params.to_h.symbolize_keys @high_similarity_matches, @low_similarity_matches, @matches = IqdbClient.new.search(**iqdb_params) diff --git a/app/logical/iqdb_client.rb b/app/logical/iqdb_client.rb index 4756336e8..03c626a8b 100644 --- a/app/logical/iqdb_client.rb +++ b/app/logical/iqdb_client.rb @@ -15,8 +15,8 @@ class IqdbClient end concerning :QueryMethods do - # Search for an image by file, URL, or post ID. - def search(post_id: nil, file: nil, url: nil, image_url: nil, file_url: nil, similarity: 0.0, high_similarity: 65.0, limit: 20) + # Search for an image by file, URL, hash, or post ID. + def search(post_id: nil, file: nil, hash: 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) @@ -31,16 +31,17 @@ class IqdbClient file = download(file_url, :image_url) elsif post_id.present? file = Post.find(post_id).file(:preview) - else - return [[], [], []] 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 } + if hash.present? + results = query_hash(hash, limit: limit) + elsif file.present? + results = query_file(file, limit: limit) + else + results = [] + end - [high_similarity_matches, low_similarity_matches, matches] + process_results(results, similarity, high_similarity) ensure file.try(:close) end @@ -59,15 +60,21 @@ class IqdbClient # Transform the JSON returned by IQDB to add the full post data for each # match. # @param matches [Array] the array of IQDB matches - # @return [Array] the array of IQDB matches, with post data - def decorate_posts(matches) + # @param low_similarity [Float] the threshold for a result to be considered low similarity + # @param high_similarity [Float] the threshold for a result to be considered high similarity + # @return [(Array, Array, Array)] the set of high similarity, low similarity, and all matches + def process_results(matches, low_similarity, high_similarity) + matches = matches.select { |result| result["score"] >= low_similarity } post_ids = matches.map { |match| match["post_id"] } posts = Post.where(id: post_ids).group_by(&:id).transform_values(&:first) - matches.map do |match| + matches = matches.map do |match| post = posts.fetch(match["post_id"], nil) match.with_indifferent_access.merge(post: post) if post end.compact + + high_similarity_matches, low_similarity_matches = matches.partition { |match| match["score"] >= high_similarity } + [high_similarity_matches, low_similarity_matches, matches] end end @@ -80,9 +87,15 @@ class IqdbClient end concerning :HttpMethods do - # Search for an image in IQDB. + # Search for an image in IQDB by hash. + # @param hash [String] the IQDB hash to search + def query_hash(hash, limit: 20) + request(:post, "query", params: { hash: hash, limit: limit }) + end + + # Search for an image file in IQDB. # @param file [File] the image to search - def query(file, limit: 20) + def query_file(file, limit: 20) media_file = MediaFile.open(file) preview = media_file.preview(Danbooru.config.small_image_width, Danbooru.config.small_image_width) file = HTTP::FormData::File.new(preview) diff --git a/app/views/iqdb_queries/show.html.erb b/app/views/iqdb_queries/show.html.erb index ecbf4eb70..00f4ade6c 100644 --- a/app/views/iqdb_queries/show.html.erb +++ b/app/views/iqdb_queries/show.html.erb @@ -10,6 +10,7 @@ <%= search_form_for(iqdb_queries_path, method: :post) do |f| %> <%= f.input :post_id, label: "Post ID", input_html: { value: params[:search][:post_id] } %> <%= f.input :url, input_html: { value: params[:search][:url] } %> + <%= f.input :hash, input_html: { value: params[:search][:hash] } %> <%= f.input :file, as: :file %> <%= f.submit "Search" %> <% end %>