Add AI tag model and UI.

Add a database model for storing AI-predicted tags, and add a UI for browsing and searching these tags.

AI tags are generated by the Danbooru Autotagger (https://github.com/danbooru/autotagger). See that
repo for details about the model.

The database schema is `ai_tags (media_asset_id integer, tag_id integer, score smallint)`. This is
designed to be as space-efficient as possible, since in production we have over 300 million
AI-generated tags (6 million images and 50 tags per post). This amounts to over 10GB in size, plus
indexes.

You can search for AI tags using e.g. `ai:scenery`. You can do `ai:scenery -scenery` to find posts
where the scenery tag is potentially missing, or `scenery -ai:scenery` to find posts that are
potentially mistagged (or more likely where the AI missed the tag).

You can browse AI tags at https://danbooru.donmai.us/ai_tags. On this page you can filter by
confidence level. You can also search unposted media assets by AI tag.

To generate tags, use the `autotag` script from the Autotagger repo, something like this:

  docker run --rm -v ~/danbooru/public/data/360x360:/images ghcr.io/danbooru/autotagger ./autotag -c -f /images | gzip > tags.csv.gz

To import tags, use the fix script in script/fixes/. Expect a Danbooru-size dataset to take
hours to days to generate tags, then 20-30 minutes to import. Currently this all has to be done by hand.
This commit is contained in:
evazion
2022-06-24 04:35:29 -05:00
parent ae9495ec7c
commit 1aeb52186e
20 changed files with 247 additions and 3 deletions

View File

@@ -0,0 +1,9 @@
<%= render(MediaAssetGalleryComponent.new(size: size)) do |gallery| %>
<% ai_tags.each do |ai_tag| %>
<% if policy(ai_tag.media_asset).can_see_image? %>
<% gallery.media_asset do %>
<%= render "ai_tags/preview", ai_tag: ai_tag, media_asset: ai_tag.media_asset, size: gallery.size %>
<% end %>
<% end %>
<% end %>
<% end %>

View File

@@ -0,0 +1,14 @@
<%= render(MediaAssetPreviewComponent.new(media_asset: media_asset, size: size, link_target: media_asset.post, html: { **data_attributes_for(media_asset) })) do |preview| %>
<% preview.footer do %>
<div class="text-center text-xs h-8">
<% if media_asset.post.present? %>
<%= link_to "post ##{media_asset.post.id}", media_asset.post %>
<% end %>
<div>
<%= link_to ai_tag.tag.pretty_name, ai_tags_path(search: { tag_name: ai_tag.tag.name, **params[:search].except(:tag_name) }), class: "tag-type-#{ai_tag.tag.category}", "data-tag-name": ai_tag.tag.name %>
<%= link_to "#{ai_tag.score}%", ai_tags_path(search: { tag_name: ai_tag.tag.name, score: ">=#{ai_tag.score}", **params[:search].except(:tag_name, :score) }), class: "tag-type-#{ai_tag.tag.category}", "data-tag-name": ai_tag.tag.name %>
</div>
</div>
<% end %>
<% end %>

View File

@@ -0,0 +1,24 @@
<%= table_for @ai_tags, class: "striped autofit" do |t| %>
<% t.column :tag do |ai_tag| %>
<%= link_to_wiki "?", ai_tag.tag.name %>
<%= link_to ai_tag.tag.pretty_name, ai_tags_path(search: { tag_name: ai_tag.tag.name }), class: "tag-type-#{ai_tag.tag.category}", "data-tag-name": ai_tag.tag.name %>
<% end %>
<% t.column :asset do |ai_tag| %>
<%= link_to "asset ##{ai_tag.media_asset_id}", ai_tag.media_asset %>
<% end %>
<% t.column :post do |ai_tag| %>
<% if ai_tag.post.present? %>
<%= link_to "post ##{ai_tag.post.id}", ai_tag.post %>
<% end %>
<% end %>
<% t.column :confidence do |ai_tag| %>
<%= ai_tag.score %>%
<% end %>
<% t.column "Present?" do |ai_tag| %>
<%= "Yes" if ai_tag.correct? %>
<% end %>
<% end %>

View File

@@ -0,0 +1,37 @@
<%= render "tags/secondary_links" %>
<div id="c-ai-tags">
<div id="a-index">
<h1>AI Tags</h1>
<%= search_form_for(ai_tags_path) do |f| %>
<%= f.input :tag_name, label: "AI Tag", input_html: { value: params.dig(:search, :tag_name).presence || Tag.find_by(id: params.dig(:search, :tag_id))&.name, data: { autocomplete: "tag" } } %>
<%= f.input :post_tags_match, label: "Post Search", input_html: { value: params.dig(:search, :post_tags_match), data: { autocomplete: "tag-query" } } %>
<%= f.input :score, label: "Confidence", input_html: { value: params.dig(:search, :score) } %>
<%= f.input :is_posted, as: :hidden, input_html: { value: params.dig(:search, :is_posted) } %>
<%= f.submit "Search" %>
<% end %>
<div class="border-b mb-4 flex flex-wrap gap-4">
<%= link_to "All", current_page_path(search: search_params.to_h.without("is_posted")), class: ["inline-block p-1 pb-2", (search_params[:is_posted].nil? ? "border-current border-b-2 -mb-px" : "inactive-link")] %>
<%= link_to "Posted", current_page_path(search: { is_posted: true, **search_params }), class: ["inline-block p-1 pb-2", (search_params[:is_posted].to_s.truthy? ? "border-current border-b-2 -mb-px" : "inactive-link")] %>
<%# link_to "Unposted", current_page_path(search: { is_posted: false, **search_params }), class: ["inline-block p-1 pb-2", (search_params[:is_posted].to_s.falsy? ? "border-current border-b-2 -mb-px" : "inactive-link")] %>
<span class="flex-grow-1"></span>
<%= render PreviewSizeMenuComponent.new(current_size: @preview_size) %>
<% if @mode == "table" %>
<%= link_to grid_icon, current_page_path(mode: nil), title: "Gallery", class: "inline-block p-1 pb-2 rounded inactive-link" %>
<% else %>
<%= link_to list_icon, current_page_path(mode: "table"), title: "Table", class: "inline-block p-1 pb-2 rounded inactive-link" %>
<% end %>
</div>
<% if params[:mode] == "table" %>
<%= render "ai_tags/table" %>
<% else %>
<%= render "ai_tags/gallery", ai_tags: @ai_tags, size: @preview_size %>
<% end %>
<%= numbered_paginator(@ai_tags) %>
</div>
</div>

View File

@@ -51,6 +51,7 @@
<li><%= link_to("Aliases", tag_aliases_path) %></li>
<li><%= link_to("Implications", tag_implications_path) %></li>
<li><%= link_to("Listing", tags_path) %></li>
<li><%= link_to("AI Tags", ai_tags_path) %></li>
<li><%= link_to("Related Tags", related_tag_path) %></li>
</ul>
<ul>

View File

@@ -4,6 +4,7 @@
<%= subnav_link_to("Aliases", tag_aliases_path) %>
<%= subnav_link_to("Implications", tag_implications_path) %>
<%= subnav_link_to "Request alias/implication", new_bulk_update_request_path %>
<%= subnav_link_to "AI tags", ai_tags_path %>
<%= subnav_link_to "Related tags", related_tag_path %>
<%= subnav_link_to "Help", wiki_page_path("help:tags") %>