diff --git a/app/controllers/mock_services_controller.rb b/app/controllers/mock_services_controller.rb index 6d1a75292..8519f670c 100644 --- a/app/controllers/mock_services_controller.rb +++ b/app/controllers/mock_services_controller.rb @@ -38,6 +38,14 @@ class MockServicesController < ApplicationController render json: @data end + def autotagger_evaluate + limit = params.fetch(:limit, 50) + predictions = tags(limit).pluck(:name).map { |name| [name, rand(0.0..1.0)] }.to_h + data = { filename: "test.jpg", tags: predictions } + + render json: data + end + private def posts(limit = 10) diff --git a/app/logical/autotagger_client.rb b/app/logical/autotagger_client.rb new file mode 100644 index 000000000..d628a74db --- /dev/null +++ b/app/logical/autotagger_client.rb @@ -0,0 +1,25 @@ +# frozen_string_literal: true + +# API client for the Danbooru Autotagger service. Used for the AI tagging feature. +# +# @see https://autotagger.donmai.us +# @see https://github.com/danbooru/autotagger +class AutotaggerClient + attr_reader :autotagger_url, :http + + def initialize(autotagger_url: Danbooru.config.autotagger_url.to_s, http: Danbooru::Http.new) + @autotagger_url = autotagger_url.chomp("/") + @http = http + end + + def evaluate(file, limit: 50, confidence: 0.01) + return {} if autotagger_url.blank? + + response = http.post("#{autotagger_url}/evaluate", form: { file: HTTP::FormData::File.new(file), threshold: confidence, format: "json" }) + return {} if !response.status.success? + + tag_names_with_scores = response.parse.first["tags"] + tags = Tag.where(name: tag_names_with_scores.keys).index_by(&:name) + tag_names_with_scores.transform_keys(&tags) + end +end diff --git a/config/danbooru_default_config.rb b/config/danbooru_default_config.rb index 7613ddd96..b8595b913 100644 --- a/config/danbooru_default_config.rb +++ b/config/danbooru_default_config.rb @@ -542,6 +542,16 @@ module Danbooru # "http://localhost:3000/mock/iqdb" end + # The URL for the Danbooru Autotagger service (https://github.com/danbooru/autotagger). Optional. + # + # Used for the AI tagging feature. Set this to http://localhost:3000/mock/autotagger + # to enable a fake server for development purposes, or do + # `docker run --rm -p 5000:5000 ghcr.io/danbooru/autotagger` to run a real server. + def autotagger_url + # "http://localhost:3000/mock/autotagger" + # "http://localhost:5000" + end + def aws_credentials Aws::Credentials.new(Danbooru.config.aws_access_key_id, Danbooru.config.aws_secret_access_key) end diff --git a/config/routes.rb b/config/routes.rb index 4a352a6fb..4827ef5d4 100644 --- a/config/routes.rb +++ b/config/routes.rb @@ -365,6 +365,7 @@ Rails.application.routes.draw do get "/mock/reportbooru/post_views/rank" => "mock_services#reportbooru_post_views", as: "mock_reportbooru_post_views" get "/mock/iqdb/query" => "mock_services#iqdb_query", as: "mock_iqdb_query" post "/mock/iqdb/query" => "mock_services#iqdb_query" + get "/mock/autotagger/evaluate" => "mock_services#autotagger_evaluate", as: "mock_autotagger_evaluate" match "*other", to: "static#not_found", via: :all end