Add IQDB for yiffy2
This commit is contained in:
parent
da068e44d0
commit
624ac2f738
2
Gemfile
2
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"
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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
|
||||
|
|
9
app/jobs/iqdb_remove_job.rb
Normal file
9
app/jobs/iqdb_remove_job.rb
Normal file
|
@ -0,0 +1,9 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
class IqdbRemoveJob < ApplicationJob
|
||||
queue_as :iqdb
|
||||
|
||||
def perform(id)
|
||||
Iqdb.remove_image(id)
|
||||
end
|
||||
end
|
11
app/jobs/iqdb_update_job.rb
Normal file
11
app/jobs/iqdb_update_job.rb
Normal 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
107
app/logical/iqdb.rb
Normal 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
|
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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 %>
|
||||
|
|
16
app/views/yiff_rest/api_v2/images/iqdb.html.erb
Normal file
16
app/views/yiff_rest/api_v2/images/iqdb.html.erb
Normal 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>
|
|
@ -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">
|
||||
|
|
|
@ -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"
|
||||
|
|
|
@ -1,4 +1,5 @@
|
|||
# frozen_string_literal: true
|
||||
|
||||
require "pagy/extras/bootstrap"
|
||||
require "pagy/extras/array"
|
||||
Pagy::DEFAULT[:items] = 20
|
||||
|
|
|
@ -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
6
db/fixes/001_iqdb_import.rb
Executable 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)
|
7
db/migrate/20241013040658_add_iqdb_id_to_api_images.rb
Normal file
7
db/migrate/20241013040658_add_iqdb_id_to_api_images.rb
Normal 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
3
db/schema.rb
generated
|
@ -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
|
||||
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue
Block a user