127 lines
3.6 KiB
Ruby
127 lines
3.6 KiB
Ruby
# 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 get_image(image, &)
|
|
Tempfile.open("yiffy2-#{image.md5}") do |tempfile|
|
|
tempfile.binmode
|
|
if image.is_external? && image.is_viewable?
|
|
HTTParty.get(image.external_api_image.file_url, Websites.config.httparty_options.merge(stream_body: true)) do |fragment|
|
|
tempfile.write(fragment)
|
|
end
|
|
tempfile.rewind
|
|
else
|
|
Websites.config.yiffy2_storage.get(image.path) do |file|
|
|
file.binmode
|
|
tempfile.write(file.read)
|
|
tempfile.rewind
|
|
end
|
|
end
|
|
yield(tempfile)
|
|
end
|
|
end
|
|
|
|
def update_image(image)
|
|
get_image(image) do |tempfile|
|
|
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
|
|
|
|
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)
|
|
get_image(image) do |tempfile|
|
|
query_file(tempfile, score_cutoff)
|
|
end
|
|
end
|
|
|
|
def query_file(file, score_cutoff)
|
|
Tempfile.open("yiffy2-queryfile") do |tempfile|
|
|
file.binmode
|
|
tempfile.binmode
|
|
tempfile.write(file.read)
|
|
tempfile.rewind
|
|
thumb = generate_thumbnail(tempfile.path)
|
|
raise(Error, "failed to generate thumb for file") unless thumb
|
|
|
|
response = make_request("/query", :post, get_channels_data(thumb))
|
|
raise(Error, "iqdb request failed") if response.status != 200
|
|
|
|
process_iqdb_result(JSON.parse(response.body), score_cutoff)
|
|
end
|
|
end
|
|
|
|
def query_hash(hash, score_cutoff)
|
|
response = make_request("/query", :post, { hash: hash })
|
|
raise(Error, "iqdb request failed") 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 || 50).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
|